ml.utils.device.gpu

GPU device type.

The default floating point type can be configured with the environment variables:

  • USE_FP64: Use FP64

  • USE_FP32: Use FP32

  • USE_BF16: Use BF16

ml.utils.device.gpu.get_env_bool(key: str) → bool[source]
class ml.utils.device.gpu.gpu_device[source]

Bases: base_device

Mixin to support single-GPU training.

classmethod has_device() → bool[source]

Detects whether or not the device is available.

Returns:

If the device is available

get_torch_compile_backend() → str | Callable[source]

Returns the backend to use for Torch compile.

Returns:

The backend

supports_grad_scaler() → bool[source]