#!/bin/bash -e

wd=$(pwd)
jobid=$(squeue --me | head -2 | tail -n1 | awk '{print $1}')


#
# Example assume allocation was created, e.g.:
# N=1 ; salloc -p standard-g  --threads-per-core 1 --exclusive -N $N --gpus $((N*8)) -t 4:00:00 --mem 0
#

module purge
module load CrayEnv
module load PrgEnv-cray/8.6.0
module load craype-x86-trento
module load craype-accel-amd-gfx90a
module load cray-python

# This path provides more recent ROCm modules.
module use module use /appl/local/containers/test-modules
module load rocm/7.0.3.lua

set -x

# Install pytorch if it doesn't exist already.
if [ ! -d $wd/pip-installs ] ; then
    pip3 install --no-cache-dir -t $wd/pip-installs --pre torch==2.10.0+rocm7.0 --index-url https://download.pytorch.org/whl/
fi

PYTHONPATH=$wd/pip-installs \
    srun --jobid=$jobid -n1 --gpus 8 \
        python -c 'import torch; print("I have this many devices:", torch.cuda.device_count())'
