#!/bin/sh
PYTHON_FILE="{YOUR PATH}/CMCS_ICLR_code/dl/main.py"
DEVICE_IDX="0"
DATASET="{CHOOSE DATASET}"
MODEL_TYPE="betatcvae"
PROJECT_NAME="{SET WANDB PROJECT NAME}"
ENTITY="{WANDB ID}"
LATENT_DIM="{SET LATENT DIMENSION}"
TRAIN_BATCH_SIZE="64"
TEST_BATCH_SIZE="512"
NUM_EPOCH="0"
MAX_STEPS="{SET TRAINING ITERATION}"
SAVE_STEPS="100000000"
SEED="1 2 3 4 5 6 7 8 9 10"
ALPHA="1.0"
BETA="{SET HYPER-PARAMETERS}"
GAMMA="1.0"
NUM_SAMPLING="10"
trap "exit" INT

for MODEL in $MODEL_TYPE
do
  for beta in $BETA
  do
    for EPOCH in $NUM_EPOCH
    do
      for BATCH in $TRAIN_BATCH_SIZE
      do
        for seed in $SEED
        do
          CUDA_VISIBLE_DEVICES=$DEVICE_IDX python $PYTHON_FILE \
          --device_idx $DEVICE_IDX \
          --dataset $DATASET\
          --project_name $PROJECT_NAME \
          --entity $ENTITY \
          --model_type $MODEL \
          --latent_dim $LATENT_DIM \
          --per_gpu_train_batch_size $BATCH \
          --test_batch_size $TEST_BATCH_SIZE \
          --num_epoch $EPOCH \
          --max_steps $MAX_STEPS \
          --save_steps $SAVE_STEPS \
          --beta $beta \
          --alpha $ALPHA \
          --gamma $GAMMA \
          --seed $seed \
          --quali_sampling $NUM_SAMPLING \
          --do_train --do_eval --do_analysis
        done
      done
    done
  done
done