lit_sb3.utils.device
[docs]
module
lit_sb3.utils.device
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34 | """PyTorch device selection utilities."""
import torch
def 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.
Args:
device: One of 'auto', 'cuda', 'mps', or 'cpu'
Returns:
Supported Pytorch device
"""
# Prefer CUDA when available, then Apple MPS.
if device == "auto":
if torch.cuda.is_available():
device = "cuda"
elif torch.backends.mps.is_available():
device = "mps"
else:
device = "cpu"
# Force conversion to torch.device
device = torch.device(device)
if device.type == "cuda" and not torch.cuda.is_available():
return torch.device("cpu")
if device.type == "mps" and not torch.backends.mps.is_available():
return torch.device("cpu")
return device
|