#!/bin/bash -l
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --gres=gpu:1
#SBATCH --time=01:00:00
#SBATCH --account=<todo>
#SBATCH --partition=<todo>
#SBATCH --output=logs/slurm_%x-%j.out

set -eo pipefail

module load ML-bundle/25.04

if [ ! -d ".venv-gh200" ]; then
    python -m venv .venv-gh200
fi
source .venv-gh200/bin/activate

export CFLAGS="${CFLAGS} -fno-lto"
export CXXFLAGS="${CXXFLAGS} -fno-lto"

pip install --upgrade pip

pip install --no-cache-dir torch==2.8.0+cu128
pip install --no-cache-dir -r requirements.txt


python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}'); print(f'Torch version: {torch.__version__}')"