Attention

Following the recent operating system upgrade, software listings and module load commands now reflect the upgraded system. Some examples and notices are still being updated and may contain outdated guidance.

PyTorch

Guidance last updated: 15 October 2026

PyTorch is an open source machine learning library for Python, used for applications such as computer vision and natural language processing.

About PyTorch on Stanage

Install PyTorch in your own environment as described under Conda on Stanage. PyTorch stopped publishing packages to its own Conda channel after version 2.5; community-maintained packages remain available on conda-forge. Official PyTorch binaries are available through pip.

The prebuilt GPU packages install their CUDA runtime and cuDNN dependencies. You do not need to load separate CUDA or cuDNN modules to run these packages. A compatible NVIDIA driver is still required. Building PyTorch or custom CUDA extensions from source requires a suitable CUDA toolkit.

GPU compatibility

Request a GPU node before running GPU calculations; see Using GPUs on Stanage. Choose a GPU-enabled PyTorch build that supports the requested GPU and its driver. The Stanage GPU types have the following CUDA compute capabilities:

GPU type

Compute capability

A100

8.0 (sm_80)

H100 and H100 NVL

9.0 (sm_90)

Use the current stable installation choices from PyTorch’s installation selector for supported stable builds.

Warning

GPU usage must be explicitly requested in submission scripts using both --partition=gpu and --gres=gpu:GPU_TYPE:N. Here, GPU_TYPE must be a100, h100 or h100nvl, and N is the number of GPUs required.

A specific GPU type should be requested because the nodes use different CPU architectures: the A100 and H100 nodes use AMD Zen 3 processors, while the H100-NVL nodes use Intel Ice Lake processors. Compiled applications and software environments should therefore be built or selected for the appropriate CPU and GPU architecture.

Using the general option --gres=gpu:1 may schedule the job on the first available GPU of any type and is not recommended unless the software has been tested on all available architectures.

To allow either an A100 or H100 GPU while excluding the H100-NVL nodes, use both --gres=gpu:1 and --exclude=gpu[31-35].

Installing PyTorch

For example, request an interactive A100 session with 4 CPUs and 16 GB of CPU RAM:

srun --partition=gpu --qos=gpu --gres=gpu:a100:1 --cpus-per-task=4 --mem=16G --time=00:45:00 --pty bash -i

Large environments can be placed in parscratch as described under Conda on Stanage.

Choose one of the following installation methods, using a separate environment for each method:

module load Anaconda3/2025.06-1
source "$EBROOTANACONDA3/etc/profile.d/conda.sh"
conda create --override-channels --channel conda-forge -n pytorch python=3.13 pytorch-gpu torchvision
conda activate pytorch

pytorch-gpu selects a GPU-enabled build. For CPU-only use, replace it with pytorch-cpu and request a CPU allocation instead.

Using PyTorch in a batch job

The following example batch submission script runs mywork.py using the Conda-forge environment created above:

#!/bin/bash
#SBATCH --partition=gpu
#SBATCH --qos=gpu
#SBATCH --gres=gpu:a100:1
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --mem=16G
#SBATCH --time=00:10:00

export SLURM_EXPORT_ENV=ALL
module load Anaconda3/2025.06-1
source "$EBROOTANACONDA3/etc/profile.d/conda.sh"
conda activate pytorch
export OMP_NUM_THREADS="$SLURM_CPUS_PER_TASK"
srun python mywork.py

For the pip installation, activate pytorch-pip instead. Load the same Anaconda module and activate your environment in each new session or job. Submit the example with sbatch my_job.slurm after saving it as my_job.slurm.

Testing your PyTorch installation

Run the following in Python to check CPU tensor operations:

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

To test GPU execution within a GPU allocation, create tensors on the CUDA device:

import torch
assert torch.cuda.is_available(), "No GPU is available to PyTorch"
x = torch.tensor([[1., 2.], [3., 4.]], device="cuda")
result = x @ x
print(result)
print(result.device)

The result should be [[7., 10.], [15., 22.]] on cuda:0. Checking torch.cuda.is_available() alone does not move your tensors or model to the GPU; use device="cuda" or .to("cuda"). See PyTorch’s CUDA semantics.

To inspect your build and GPU, use:

print(torch.__version__)
print(torch.version.cuda)
print(torch.cuda.get_device_name(0))
print(torch.cuda.get_device_capability(0))
print(torch.cuda.get_arch_list())