PyTorch

PyTorch is a commonly used Python package for deep learning.

Basic usage

First, check the tutorials up to and including GPU computing.

If you plan on using NVIDIA’s containers to run your model, please check the page about NVIDIA’s singularity containers.

The basic way to use PyTorch is via the Python in the scicomp-python-env module. Don’t load any additional CUDA modules, scicomp-python-env includes everything.

On newer GPUs like B300, you will need to load the new Pytorch module scicomp-pytorch-env/2026.1 to get the latest version of PyTorch that supports these GPUs. However, this module does not work on older GPUs like the V100s due to differences in CUDA compute capability requirements and supported GPU architectures.

Building your own environment with PyTorch

If you need a PyTorch version different to the one supplied with the scicomp python environment we recommend installing your own conda environment as detailed here.

Creating an environment with GPU enabled PyTorch

To create an environment with GPU enabled PyTorch you can use an environment file like this pytorch-env.yml:

name: pytorch-env
channels:
  - nvidia
  - conda-forge
dependencies:
  - python==3.12
  - pip

  # PyTorch ecosystem
  - pytorch
  - pytorch-gpu
  - torchmetrics
  - torchaudio
  - torchvision
  - torchvision-extra-decoders

Here we install the latest gpu enabled pytorch version from conda-forge-channel.

Hint

During installation conda will try to verify what is the maximum version of CUDA installed graphics cards can support and it will install non-CUDA enabled versions by default if none are found (as is the case on the login node, where environments are normally built). This can be usually overcome by setting explicitly that the packages should be the CUDA-enabled ones. It might however happen, that the environment creation process aborts with a message similar to:

nothing provides __cuda needed by tensorflow-2.9.1-cuda112py310he87a039_0

In this instance it might be necessary to override the CUDA settings used by conda/mamba. To do this, prefix your environment creation command with CONDA_OVERRIDE_CUDA=CUDAVERSION, where CUDAVERSION is the CUDA toolkit version you intend to use as in:

CONDA_OVERRIDE_CUDA="11.2" mamba env create -f cuda-env.yml

This will allow conda to assume that the respective CUDA libraries will be present at a later point and so it will skip those requirements during installation.

For more information, see this helpful post in Conda-Forge’s documentation.

Examples:

Let’s run the MNIST example from PyTorch’s tutorials:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

The full code for the example is in pytorch_mnist.py. One can run this example with srun:

$ wget https://raw.githubusercontent.com/AaltoSciComp/scicomp-docs/master/triton/examples/pytorch/pytorch_mnist.py
$ module load scicomp-python-env
$ srun --time=00:15:00 --mem=4G --gpus=1 python pytorch_mnist.py

or with sbatch by submitting pytorch_mnist.sh:

#!/bin/bash
#SBATCH --time=00:15:00
#SBATCH --mem=4G
#SBATCH --gpus=1

module load scicomp-python-env

python pytorch_mnist.py

The new module scicomp-pytorch-env/2026.1 can also be used to run the example, but you will need to specify the min CUDA version in your batch script.

pytorch_mnist_min_cuda.sh:

#!/bin/bash
#SBATCH --time=00:15:00
#SBATCH --mem=4G
#SBATCH --gpus=1
#SBATCH --gres=min-cuda-cc:80

module load scicomp-pytorch-env/2026.1

python pytorch_mnist.py

The Python-script will download the MNIST dataset to data folder.

Let’s run the MNIST example from PyTorch’s tutorials:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

The full code for the example is in pytorch_mnist.py. One can run this example with srun:

wget https://raw.githubusercontent.com/AaltoSciComp/scicomp-docs/master/triton/examples/pytorch/pytorch_mnist.py
module load nvidia-pytorch/20.02-py3
srun --time=00:15:00 --mem=4G --gpus=1 singularity_wrapper exec python pytorch_mnist.py

or with sbatch by submitting pytorch_singularity_mnist.sh:

#!/bin/bash
#SBATCH --time=00:15:00
#SBATCH --mem=4G
#SBATCH --gpus=1

module load nvidia-pytorch/20.02-py3

singularity_wrapper exec python pytorch_mnist.py

The Python-script will download the MNIST dataset to data folder.