composer.trainer.devices.device_gpu#

The GPU device used for training.

Classes

DeviceGPU

An extension of Device for GPUs.
class composer.trainer.devices.device_gpu.DeviceGPU(device_id=None, *, allow_tf32=True)[source]#

Bases: composer.trainer.devices.device.Device

An extension of Device for GPUs.

Parameters

  • device_id (int, optional) โ€“ Integer ID of a GPU device to train with. If not specified, the local rank of the current process is used. Default: None.

  • allow_tf32 (bool, optional) โ€“ Whether to allow TF32 matrix multiplications. Defaults to True. For more information, see TensorFloat-32(TF32) on Ampere devices.