#!/bin/bash
#SBATCH --job-name=torch_dl_job
#SBATCH --partition=a40             # Target partition for the A40 nodes
#SBATCH --output=pytorch_job_%j.out
#SBATCH --error=pytorch_job_%j.err
#SBATCH --nodes=1                   # Request resources on a single node
#SBATCH --ntasks=1                  # Run a single task instance
#SBATCH --gres=gpu:1                # Request 1 A40 GPU allocation for tensor calculations
#SBATCH --cpus-per-task=4           # Request 4 CPU worker threads for background data loading
#SBATCH --time=01:00:00             # Max job runtime allocation (HH:MM:SS)

# Clean out inherited environment parameters
module purge

# Crucial: Load the exact same Python and PyTorch module pair
module load unite/python/3.14/python-3.14.0
module load unite/python/3.14/pytorch

# Run the PyTorch training or processing script
python3 pytorch_test.py