torchtnt.utils.env.init_from_env

torchtnt.utils.env.init_from_env(*, device_type: Optional[str] = None, dist_init_method_type: Literal['env', 'tcp', 'file'] = 'env', pg_backend: Optional[str] = None, pg_timeout: timedelta = datetime.timedelta(seconds=1800), float32_matmul_precision: str = 'high') → device

Utility function that initializes the device and process group, if applicable.

The global process group is initialized only if:

• torch.distributed is available is not already initialized
• the program has been launched on multiple processes

This is intended as a convenience to include at the beginning of scripts that follow a SPMD-style execution model.

Parameters:

• device_type (str, optional) – Device type to initialize. If None, device will be initialized based on environment
• dist_init_method_type (str, optional) – Method to initialize the process group. Must be one of “env”, “tcp”, or “file”. For more information, see here: https://pytorch.org/docs/stable/distributed.html#initialization
• pg_backend (str, optional) – The process group backend to use. If None, it will use the default process group backend from the device
• pg_timeout (timedelta, optional) – Timeout for operations executed against the process group. Default value equals 30 minutes
• float32_matmul_precision (str, optional) – The setting for torch’s precision of matrix multiplications.

Returns:

The current device.