device module

Documentation for hippynn.experiment.device module.

This file is used to host the function set_devices, as it is needed by both serialization and routines. We may find a better place for this function in the future.

set_devices(model: GraphModule, loss: GraphModule, evaluator: Evaluator, optimizer: Optimizer, device: device | str | Tuple | list) → Tuple[source]

Sets the model, loss, and optimizer to the specified device if necessary. Evaluation loss is performed on CPU.

Parameters:

  • model (GraphModule) – current model on CPU

  • loss (GraphModule) – current loss module on CPU

  • evaluator (Evaluator) – evaluator

  • optimizer (torch.optim.Optimizer) – optimizer with state dictionary on CPU

  • device (Union[torch.devices, string, tuple, list]) – target device for training. If device is a tuple or list, wrap the model using torch.nn.DataParallel and use the first specified device as the “primary: one.

Returns:

model, evaluator, optimizer

Return type:

Tuple