Skip to content

PyTorch

PyTorch is an optimized tensor library for deep learning using GPUs and CPUs.

Homepage Version(s): 2.9.1

Installing PyTorch on the HPC#

PyTorch is not provided through the system-wide Python installation on the HPC. Users should install PyTorch in a user-managed Python environment. We recommend using either uv (a fast Python package and environment manager) or a Python virtual environment (venv).

Installing PyTorch with uv#

uv is a fast, modern Python package and environment manager that replaces tools such as pip, and virtualenv. First, load the uv module: 

1
$ module load python-uv

Then create and activate a uv environment: 

1
2
3
4
$ mkdir pytorch_env
$ cd pytorch_env
$ uv venv
$ source .venv/bin/activate

Next, install PyTorch using uv: 

1
(pytorch_env) $ uv pip install torch torchvision torchaudio

Installing PyTorch with a Python Virtual Environment (venv)#

First, load the Python module: 

1
$ module load python

Then create and Activate a virtual environment: 

1
2
$ python -m venv ~/venvs/pytorch_env
$ source ~/venvs/pytorch_env/bin/activate

Next, install PyTorch using pip: 

1
(pytorch_env) $ pip install torch torchvision torchaudio

Using PyTorch on the HPC#

Use the sample code from the Verification section of the PyTorch guide:

1
2
3
import torch
x = torch.rand(5, 3)
print(x)

Using GPUs with PyTorch#

To run a PyTorch job on a GPU node, you need to submit a job to our job scheduler, and indicate that you want to run using GPUs.

Warning

Running torch.cuda.is_available() on the login node will return False. It will return True on a GPU node.

Example Slurm GPU Submit script#

1
2
3
4
5
6
7
8
9
#!/bin/bash

#SBATCH -A backfill2   # Submit to backfill2 (note, this account has a four-hour walltime limit)
#SBATCH -n 12          # Request resources for 12 tasks
#SBATCH --gres=gpu:1   # Request a single GPU card

conda activate pytorch_app  # Substitute your conda environment for 'pytorch_app'

srun python MYCODE.py