#!/usr/bin/env bash
set -euo pipefail

DATA=${1:?usage: $0 /path/to/dataset_dir}
SAVE_DIR="${SAVE_DIR:-./runs/fewshot_rtm}"

mkdir -p "$SAVE_DIR"

python train.py \
    --hps fewshot \
    --save_dir "$SAVE_DIR" \
    --data_root "$DATA" \
    --num_epochs 15000 \
    --fid_freq 1000 \
    --latent_dim 1024 \
    --eps_radius 0.12 \
    --use_rtm True \
    --H_cycles 4 --L_cycles 2 --refinement_steps 2 \
    --num_tokens 4 \
    --rtm_hidden_size 256