Software Frameworks

PyTorch Setup

Prerequisites: (Make sure to read the following before using this example!)

#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=1
#SBATCH --mem=16G
#SBATCH --time=00:15:00
#SBATCH --partition=unkillable

set -e  # exit on error.
echo "Date:     $(date)"
echo "Hostname: $(hostname)"

module purge
# This example uses Conda to manage package dependencies.
# See for more information.
module load anaconda/3

# Creating the environment for the first time:
# conda create -y -n pytorch python=3.9 pytorch torchvision torchaudio \
#     pytorch-cuda=11.6 -c pytorch -c nvidia

# Activate the environment:
conda activate pytorch


import torch
import torch.backends.cuda

def main():
    cuda_built = torch.backends.cuda.is_built()
    cuda_avail = torch.cuda.is_available()
    device_count = torch.cuda.device_count()

    print(f"PyTorch built with CUDA:         {cuda_built}")
    print(f"PyTorch detects CUDA available:  {cuda_avail}")
    print(f"PyTorch-detected #GPUs:          {device_count}")
    if device_count == 0:
        print("    No GPU detected, not printing devices' names.")
        for i in range(device_count):
            print(f"    GPU {i}:      {torch.cuda.get_device_name(i)}")

if __name__ == "__main__":

Running this example

$ sbatch