For our CPU systems, 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 ensure a cpu-only installation, you can specify that with something like: 

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

 Basic PyTorch Test

The following command can be run in an allocation to test your environment:

python -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())'

The output should look like:

tensor([[0.7898, 0.4950, 0.3293],
        [0.2917, 0.5806, 0.6846],
        [0.3356, 0.9470, 0.3803],
        [0.5355, 0.6018, 0.6081],
        [0.8721, 0.1696, 0.0208]])
Torch Version 2.4.0+cu121
Built for Cuda 12.1
Cuda available: False

Launching a Jupyter Notebook

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

python -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.