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

source ~/pytorch-2.1.0/bin/activate
# Where ~/pytorch-2.1.0/ is the path to the virtual environment
python3 pytorch_test.py