DeviceGPU#

class composer.devices.DeviceGPU(device_id=None, *, allow_tf32=True)[source]#

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.