PyTorch

Published

January 18, 2026

This page shows an example of using PyTorch. PyTorch is a popular deep learning library for training artificial neural networks.

It can be installed in many ways. Optimized versions with Easybuild are available in the software stack on ADA. Search the module environment for the appropiate version.

Example GPU job

This example recipe from the PyTorch tutorials site measures the performance of a simple network in default precision, then walks through adding autocast and GradScaler to run the same network in mixed precision with improved performance.

It can be run as a standalone python script. Download the script:

wget https://raw.githubusercontent.com/pytorch/tutorials/refs/heads/main/recipes_source/recipes/amp_recipe.py

This script can be run interactively

python3 amp_recipe.py

But we will create a slurm script Torch-AMP-ex.sh to run on a compute node in batch mode. In this example 1 Node, 2 cores and 1 GPU are requested for 10 minutes.

#!/bin/bash -l
#SBATCH -J Torch-AMP-example
#SBATCH -N 1
#SBATCH --ntasks-per-node=2
#SBATCH --gpus=1
#SBATCH --time=0-00:10:00
#SBATCH --mail-type=end,fail

echo "== Starting run at $(date)"
echo "== Job ID: ${SLURM_JOBID}"
echo "== Node list: ${SLURM_NODELIST}"
echo "== Submit dir. : ${SLURM_SUBMIT_DIR}"
echo "== Scratch dir. : ${TMPDIR}"

# 1. Load the environment
module load 2024
module load PyTorch/2.1.2-foss-2023a-CUDA-12.1.1

# 2. (Optional) Install missing dependency seen in logs
pip install --user tqdm > /dev/null 2>&1 || true

# 3. Copy the python script to scratch so we don't modify the original
cp amp_recipe.py ${TMPDIR}/amp_recipe_patched.py

# 4. Patch the script compatibility issues using sed
sed -i 's/torch\.amp\.GradScaler("cuda")/torch.cuda.amp.GradScaler()/g' ${TMPDIR}/amp_recipe_patched.py
sed -i 's/torch\.amp\.GradScaler("cuda"\s*,/torch.cuda.amp.GradScaler(/g' ${TMPDIR}/amp_recipe_patched.py

# 5. Run the patched script from scratch
echo "== Running patched script..."
python3 ${TMPDIR}/amp_recipe_patched.py

You can monitor the status of the job with squeue -u $USER. Once the job runs, you’ll have a slurm-xxxxx.out file in the directory. This log file contains both PyTorch and Slurm output.

Performance and Results

Depending on the type of GPU you may find different performance.For example:

GPU	NVIDIA RTX 2070 super	NVIDIA A30
Default precision:	6.835 sec	4.444 sec
Mixed precision:	4.747 sec	0.858 sec