Pytorch Cookbook
Auto Set Device
import torch
print(f"Default device: {torch.get_default_device()}")
if torch.backends.mps.is_available():
device = torch.device("mps")
x = torch.ones(1, device=device)
print(f"Using MPS device: {x}")
elif torch.backends.cuda.is_built():
device = torch.device("cuda")
x = torch.ones(1, device=device)
print(f"Using CUDA device: {x}")
else:
device = torch.device("cpu")
x = torch.ones(1, device=device)
print(f"Using CPU device: {x}")
device = torch.device((
"cuda"
if torch.cuda.is_available()
else "mps"
if torch.backends.mps.is_available()
else "cpu"
))
print(f"Using {device} device")
# 不推荐
torch.set_default_device(device)