#!/bin/bash
#SBATCH --job-name=pytorch-case
#SBATCH --partition=ampere
#SBATCH --gpus=1
#SBATCH --ntasks=1
#SBATCH --time=10:00

module load gcc/13.2.0-i python/3.11 cuda/12.4.0
source ~/pytorch-env/bin/activate
# Where ~/pytorch-2.1.0/ is the path to the virtual environment
python3 pytorch_test.py