#!/bin/bash -l
#SBATCH --time=24:00:00
#SBATCH -J afno
#SBATCH -o afno_backbone_finetune.out
#SBATCH -N 1 -n 1 -c 8 --gres=gpu:a100:1 -p GPU-8A100 --qos=gpu_8a100

config_file=./AFNO.yaml
config='afno_backbone_ustc'
run_num='d6p2_lr2'

export HDF5_USE_FILE_LOCKING=FALSE
export NCCL_NET_GDR_LEVEL=PHB

export MASTER_ADDR=$(hostname)

set -x
srun -u --mpi=pmi2 \
    bash -c "
    source /home/ess/cxt/miniconda3/etc/profile.d/conda.sh    
    conda activate pytorch
    python ../train.py --enable_amp --yaml_config=$config_file --config=$config --run_num=$run_num
    "