PyTorch

Setup Mamba/Conda Environment

Setup and activate a new Mamba environment as outlined in the Python Documentation. Install the following packages inside your mamba environment:

mamba install -c nvidia -c pytorch cudatoolkit=11.8.0 pytorch

To install a specific version, you can run:

mamba install -c nvidia -c pytorch cudatoolkit=X.Y.Z pytorch=X.Y.Z

Job Script Specifications

• To use a Conda environment inside of a Slurm job script, modify the first line so it looks like this:

#!/bin/bash --login

• Request a GPU in your script:

#SBATCH --gpus=1 # NUMBER OF GPU'S

• PyTorch doesn't like to play nice with the Kepler (K80) nodes. A good alternative is the Pascal (P100) nodes. There are other capable GPU's, but the Pascal is the most common. To use them, add the following to your jobscript:

#SBATCH --constraint='pascal' # SPECIFY PASCAL NODES
# # To use other GPU's, substitute 'pascal' in the constraint

• Load Mamba Environment At The Bottom Of The Job Script:

# # LOAD MODULES, INSERT CODE, AND RUN YOUR PROGRAMS HERE
mamba activate <EnvironmentName>
# # RUN ANY SCRIPTS, CODE, OR PROGRAMS THAT RELY ON THE MAMBA ENVIRONMENT HERE