Shortcuts

get_device

torchtune.utils.get_device(device: Optional[str] = None) device[source]

Function that takes an optional device string, verifies it’s correct and available given the machine and distributed settings, and returns a torch.device. If device string is not provided, this function will infer the device based on the environment.

If CUDA is available and being used, this function also sets the CUDA device.

Parameters:

device (Optional[str]) – The name of the device to use.

Returns:

device.

Return type:

torch.device

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources