#!/bin/bash
#SBATCH --job-name=Nefeli-pytorch
#SBATCH --nodes=1
#SBATCH --partition=a100
#SBATCH --gpus=2
#SBATCH --time=10:00

module load gcc/14.2.0 python/3.14.0 cuda/12.8.1
source ~/pytorch-env/bin/activate
# Where ~/pytorch-2.1.0/ is the path to the virtual environment
python3 pytorch_test.py