Installing PyTorch with GPU support using Conda.

The easiest way to install PyTorch with GPU support is using Conda. This will yield a Python environment, located in your $WORK directory, with a GPU-enabled version of the PyTorch package.

First, you need to open a terminal connection to a node with a GPU. The following command will do this:

srun --partition=2080-galvani --gres=gpu:1 --pty bash

Now you are in a terminal window on a compute node. On the new terminal on the compute node, run the following commands. First, you are creating a new Conda environment with Python version 3.11. Second, you are activating that environment so that you can run commands within it. Third, you are installing the PyTorch package with CUDA/GPU support.

conda create -p $WORK/.conda/py-311-pytorch python=3.11
conda activate $WORK/.conda/py-311-pytorch
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

If you need other Python packages in your PyTorch environment, you can install them with python3 -m pip install $PACKAGE_NAME or conda install $PACKAGE_NAME. To find which packages are available, you must search the PyPi website.

If you want to check that this worked, just after running the three commands above, run these two commands at the terminal.

python3 -c "import torch; print(torch.cuda.is_available())"
# Should return True


python3 -c "import torch; print(torch.rand(5, 3))"
# Should return tensor([…]) with 3 columns and 5 rows with random numbers inside.

Using PyTorch with GPU support

To use PyTorch with GPU support, run the following in your sbatch script:

conda activate $WORK/.conda/py-311-pytorch

This will activate the Conda environment for PyTorch (plus whatever other packages you have in that environment).