lit_sb3.utils.device

source module lit_sb3.utils.device

PyTorch device selection utilities.

Functions

• get_device — Retrieve PyTorch device.

source get_device(device: torch.device | str = 'auto') → torch.device

Retrieve PyTorch device.

It checks that the requested device is available first. By default, it prefers CUDA, then MPS, then CPU.

Parameters

• device : torch.device | str — One of 'auto', 'cuda', 'mps', or 'cpu'

Returns

• torch.device — Supported Pytorch device