| 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 | |