Spaces:
Running
Running
File size: 599 Bytes
96170c6 2207aa9 96170c6 2207aa9 96170c6 2207aa9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | import torch
def get_device(device=None):
"""
Pick the best available device.
Priority: explicit override -> CUDA (Colab / lab box) -> MPS (Apple Silicon)
-> CPU. Forward-only eval on MPS is solid in torch >= 2.
Args:
device: Manual override ('cuda', 'mps', 'cpu', or None for auto)
Returns:
torch.device
"""
if device is not None:
return torch.device(device)
if torch.cuda.is_available():
return torch.device("cuda")
if torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
|