lmflow.utils.envs

ref: https://github.com/pytorch/torchtune/blob/main/torchtune/utils/_device.py

Attributes

logger
is_cuda_available

Functions

is_accelerate_env()
get_device_name(→ str)	Get the device name based on the current machine.
get_torch_device(→ any)	Return the corresponding torch attribute based on the device type string.

Module Contents

lmflow.utils.envs.logger

lmflow.utils.envs.is_cuda_available

lmflow.utils.envs.is_accelerate_env()

lmflow.utils.envs.get_device_name() → str

Get the device name based on the current machine.

lmflow.utils.envs.get_torch_device() → any

Return the corresponding torch attribute based on the device type string. Returns:

module: The corresponding torch device namespace, or torch.cuda if not found.