Device

class composer.devices.Device

Abstract class for a device on which a model runs.

dist_backend

Distributed backend to use. Should be gloo, mpi, or nccl. See the pytorch docs for details.

Type

str

batch_to_device(batch)

Invoked by the Trainer move all tensors items in a batch to device.

Supports nested sequences and mappings of tensors. Ignores non-tensor items. Preserves sequence and mapping types when possible; otherwise, sequences are converted to lists, and mappings are converted to dictionaries.

Parameters

batch (Any) – The batch to move to the device.

Returns

Batch – The batch on the device.

abstract module_to_device(module)

Invoked by the Trainer to move a module onto the device.

Parameters

module (Module) – The module to move to the device.

Returns

torch.nn.Module – The module on the device.

optimizer_to_device(optimizer)

Invoked by the Trainer to move the optimizer’s state onto the device.

Parameters

optimizer (Optimizer) – The optimizer to move to the device

Returns

Optimizer – The optimizer on the device

abstract tensor_to_device(tensor)

Invoked by the Trainer to move a tensor onto a device.

Parameters

tensor (Tensor) – The tensor to move to the device.

Returns

Tensor – The tensor on the device.