code/train_rtm.slurm

#!/bin/bash
#SBATCH -A mhahsler_course_recomm_0001
#SBATCH -N 1
#SBATCH -G 1
#SBATCH -c 8
#SBATCH --mem=64G
#SBATCH --time=1-0:00:00
#SBATCH --requeue
#SBATCH -J rtm_cifar10
#SBATCH --chdir=/work/projects/mhahsler/course_recomm/allocation001/AI_Club/paper/TRMisDiffusion/RTM_latent_refinement
#SBATCH -o logs/rtm_cifar10_%j.out
#SBATCH -e logs/rtm_cifar10_%j.err
#SBATCH --mail-user=jerryma@smu.edu
#SBATCH --mail-type=BEGIN,END,FAIL
#SBATCH --signal=USR1@120

# ── Environment ────────────────────────────────────────────────────────────
module load gcc/13.2.0
module load cuda/12.8.0-6mc4fti


if [ ! -d "venv" ]; then
    virtualenv -p python venv
    source venv/bin/activate
    pip install -r requirements.txt
    pip install -i https://test.pypi.org/simple/ dciknn-cuda==0.1.15
    pip install wandb huggingface_hub transformers faiss-gpu imageio clean-fid
else
    source venv/bin/activate
fi

if [ ! -d "datasets/cifar10" ]; then
    python prepare_cifar10.py --data_root ./datasets/cifar10
fi

# Run training
NUM_GPUS=1 bash scripts/train_cifar10.sh

# Run HuggingFace upload script
python upload_to_hf.py