#!/bin/bash
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --gpus=4
#SBATCH --cpus-per-task=18
#SBATCH -t 4:00:00
#SBATCH --output=outputs/logs/cifar10_neoneural_cat_mlp
#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 data=cifar10 batch_size=256 data.train.drop_rate=0.0 \
#   data.train.noise_scale=0.05 mixup=True num_epochs=100 optim.lr=0.0005 model=cat_mlp \
#   reuse_backbone=False model.hidden_dim=1024 optim.weight_decay=0.1 model.dropout=0.3 \
#   model.num_embeddings=6 eval_every=10000 $EXTRA_ARGS

srun python -u main.py data=cifar10 batch_size=256 data.train.drop_rate=0.0 \
  data.train.noise_scale=0.05 mixup=True num_epochs=100 optim.lr=0.0008 model=cat_mlp \
  reuse_backbone=False model.hidden_dim=2048 optim.weight_decay=0.05 model.dropout=0.3 \
  model.num_embeddings=6 eval_every=10000 $EXTRA_ARGS
