#!/usr/bin/env bash

# Slurm job configuration
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --ntasks-per-node=1
#SBATCH --output=output_%j.out
#SBATCH --error=error_%j.er
#SBATCH --time=00:10:00
#SBATCH --job-name=PYTORCH_MNIST
#SBATCH --gres=gpu:1 --partition=develgpus
#SBATCH --mail-type=ALL

# Load the required modules
module use /usr/local/software/jureca/OtherStages
module load Stages/2018b
module load GCC/7.3.0
module load MVAPICH2/2.3-GDR
module load PyTorch/1.0.0-GPU-Python-3.6.6
module load torchvision/0.2.1-GPU-Python-3.6.6

# Run the program
srun python -u mnist.py