DeviceGPU

class composer.devices.DeviceGPU(device_id=None, *, allow_tf32=True)

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 (and later) devices.