composer.trainer.devices.device_gpu#
The GPU device used for training.
Classes
- 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.