For our systems using NVIDIA GPUs (Matrix and RZVector), we recommend you install PyTorch using a cached wheel (https://wci-repo.llnl.gov/#browse/browse:pypi-wci:torch).

Option 1: Using a System Site PyTorch

To use a system-provided PyTorch installation, users must start with one of the python modules we list above. We then recommend that users create a virtual environment. The name "mytorchenv" is a variable that can be changed to meet your own needs.

module load python/3.12
python3 -m venv --system-site-packages mytorchenv
source mytorchenv/bin/activate

Option 2: Installing Your Own PyTorch

To install your own PyTorch, you should create a virtual environment without system site packages. The name "mytorchenv" is a variable that can be changed to meet your own needs.

module load python/3.12
python3 -m venv mytorchenv
source mytorchenv/bin/activate
pip install torch torchvision torchaudio 

This pip install command should fetch a default torch (and dependencies) version from Nexus. 

If you need to explicitly control the torch version or cuda version, you can specify that with something like: 

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

 Basic PyTorch Test

Run the following command in a GPU allocation 

python3 -c 'import torch ; print(torch.rand(5, 3)) ; print("Torch Version", torch.__version__) ; print("Built for Cuda", torch.version.cuda) ; print("Cuda available:", torch.cuda.is_available())'

Launching a Jupyter Notebook

To make your virtual environment (with torch installed) available via a Jupyter Notebook, you must run the following command:

python3 -m ipykernel install --prefix=$HOME/.local/ --name 'mytorchkernel' --display-name 'My Torch kernel'

The name and display name arguments can be customized to meet your own needs. For most personal-use installations, these are the only commands you should need to run to make the torch kernel available to a Jupyter Notebook.

Additional details and setup information can be found on the Orbit and Jupyter Notebooks page.