PyTorch on Apple Silicon

Already some time ago, PyTorch became fully available for Apple Silicon. Thus, it’s no longer necessary to install the nightly builds to run PyTorch on the GPU of your Apple Silicon machine as I described in one of my earlier posts.

Install PyTorch on Apple Silicon

1. Let’s begin with creating a new conda environment:

     conda create -n pytorch_env -y python=3.12
     conda activate pytorch_env
     conda install -y mamba
   

2. Now, we can install PyTorch either via conda/mamba or pip:

     mamba install pytorch torchvision torchaudio -c pytorch
   

   or

     pip install torch torchvision torchaudio
   

That’s it. You can now start using PyTorch on your Apple Silicon machine. To verify the installation, you can run the following code:

import torch
print(f"torch backend MPS is available? {torch.backends.mps.is_available()}")
print(f"current PyTorch installation built with MPS activated? {torch.backends.mps.is_built()}")
print(f"check the torch MPS backend: {torch.device('mps')}")
print(f"test torch tensor on MPS: {torch.tensor([1,2,3], device='mps')}")

Don’t forget to move your model to the MPS device, if you want to use it:

device = torch.device('mps')
model = model.to(device)

You can monitor the GPU usage of your Apple Silicon machine using the Activity Monitor (View->GPU History). This helps you to verify that your model is indeed running on the GPU.

Conclusion

I tested the installation on my M1 Max MacBook Pro with macOS Sonoma 14.3.1 using Python 3.8 to 3.12, and it worked without any issues, both for the installation via conda and pip. If you encounter any issues, feel free to leave a comment below.