Skip to content

Instantly share code, notes, and snippets.

@iamaziz

iamaziz/pytorch_device.py

Created December 30, 2023 06:57
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save iamaziz/ec9a48c07deb06000820233651235090 to your computer and use it in GitHub Desktop.
Save iamaziz/ec9a48c07deb06000820233651235090 to your computer and use it in GitHub Desktop.
Download ZIP
Enable MacOS M-series chip for PyTorch
Raw
pytorch_device.py
import torch
# Check if CUDA is available, else check for MPS, otherwise default to CPU
if torch.cuda.is_available():
device = torch.device("cuda") # GPU
elif torch.backends.mps.is_available():
device = torch.device("mps") # MacOS M-series chip
else:
device = torch.device("cpu") # CPU
print(f"Using device: {device}")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment