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

# 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_h5.py data=shapenet num_epochs=20 num_finetune_epochs=1 save_ckpt=True \
  model.embed_dim=512 model.num_nodes=16 model.pos_embedding_dim=512 \
  model.neomlp_attention.ffn_dim=512 train_fraction=0.1 finetune_fraction=0.2 \
  num_train_images=32768 num_eval_images=1 steps_till_eval=2000 $EXTRA_ARGS

