Torch - Get device across CPU / GPU / MPS

Posted on June 15, 2024 • Tags: python torch mps cuda ml

To check if you can run on CUDA (Nvidia GPU), MPS (Mac), or CPU:

import torch

device: str = 'cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu')

Helpful PDB Commands >>