#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --gpus=1
#SBATCH --cpus-per-task=18
#SBATCH -t 3-0
#SBATCH --output=outputs/logs/cifar10_neomlp
#SBATCH --open-mode=append
#SBATCH --partition=gpu

# module load 2023
# module load Mamba/23.1.0-4

source $HOME/.bashrc
mamba activate neomlp

echo "=================================="
echo "Date: $(date)"
echo "Directory: $(pwd)"
python -c 'import torch; print(torch.randn(1).cuda())'

EXTRA_ARGS=$@
echo "Extra args: $EXTRA_ARGS"
# cd experiments/neural_datasets
echo `ls`
srun python -u main.py num_epochs=20 num_finetune_epochs=5 save_ckpt=True \
  num_train_images=16 num_eval_images=16 model.embed_dim=512 model.num_nodes=8 \
  model.pos_embedding_dim=512 model.neomlp_attention.ffn_dim=512 num_augmentations=50 \
  steps_till_eval=3125 $EXTRA_ARGS

