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


module load 2022
module load Anaconda3/2022.05
# module load CUDA/11.3.1

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/high_res_inr
echo `ls`
srun python -u main.py data=bigbuckbunny data.batch_size=4096 data.eval_batch_size=4096 \
  data.steps_till_eval=50000 optim.lr=5e-4 model.neomlp_attention.ffn_dim=1024 \
  model.embed_dim=256 model.neomlp_attention.num_heads=8 model.num_nodes=16 \
  model.neomlp_attention.num_layers=4 model.pos_embedding_dim=128 $EXTRA_ARGS

