PyTorch

Setup Mamba/Conda Environment

Setup and activate a new mamba environment as outlined on the Conda Environments page.

To install PyTorch, we need to select the right version. Running

conda search pytorch

will result in a long list similar to the following:

# Name                       Version           Build  Channel             
...
pytorch                        2.7.1 cpu_mkl_py313_he78a34b_102  conda-forge         
pytorch                        2.7.1 cpu_mkl_py313_hea9ba1b_100  conda-forge         
pytorch                        2.7.1 cpu_mkl_py39_h08fd187_101  conda-forge         
pytorch                        2.7.1 cpu_mkl_py39_h08fd187_102  conda-forge         
pytorch                        2.7.1 cpu_mkl_py39_h4c5010e_103  conda-forge         
pytorch                        2.7.1 cpu_mkl_py39_h92ffbb0_100  conda-forge         
pytorch                        2.7.1 cuda126_generic_py310_hd77d226_200  conda-forge         
pytorch                        2.7.1 cuda126_generic_py311_h170ca9b_200  conda-forge         
pytorch                        2.7.1 cuda126_generic_py312_h261cdbc_200  conda-forge         
pytorch                        2.7.1 cuda126_generic_py313_h14c909a_200  conda-forge         
pytorch                        2.7.1 cuda126_generic_py39_h98ce218_200  conda-forge         
pytorch                        2.7.1 cuda126_mkl_py310_h5ee0071_300  conda-forge         
pytorch                        2.7.1 cuda126_mkl_py311_hcada2b2_300  conda-forge         
pytorch                        2.7.1 cuda126_mkl_py312_h30b5a27_300  conda-forge         
pytorch                        2.7.1 cuda126_mkl_py313_he20fe19_300  conda-forge         
pytorch                        2.7.1 cuda126_mkl_py39_hd241233_300  conda-forge  

Note: Running mamba search will result in similar builds being consolidated. This is fine for most installations, but not when we need this level of detailed information about PyTorch.

Conda/mamba will default to installing the latest version PyTorch, regardless of if we support the CUDA version or it was only built for CPU runs. Therefore, we need to specify which build we want to install. It's always recommended to install the latest version your code can run. For this example, we're going to use 2.7.1.

Next, we'll look at the build category. Each of the strings start with either "cpu" or "cudaXXX". This means the build can only run on the CPU or GPU respectively. Please choose the GPU version. In order to know which CUDA our drivers' support, run nvidia-smi and look at the top right hand corner of the output. It will tell you the latest version of CUDA we support. This number gets updated regularly. Choose "cudaXXX" where "XXX" is less than or equal to the latest version we support.

Since most of our GPUs are paired with Intel CPUs, we'll choose "mkl" (Intel's Math Kernel Libraray) instead of "generic". This shouldn't matter too much if you are using the GPU properly. A CPU installation should choose "generic."

Finally, we need to select the Python version. Usually, the latest Python is best, but choose one that will work with your code. If you already have Python installed in your environment, match that version.

The hash at the end provides no further information.

Putting this all together, we would install PyTorch like so

mamba install pytorch=2.7.1=cuda126_mkl_py313_he20fe19_300

Job Script Specifications

Make sure to add the --gpus=N flag to your job script where N is the number of GPUs you need. Additionally, remember to load miniforge3 and activate your environment before running code.