export CUDA_VISIBLE_DEVICES=4,5
torchrun --nproc_per_node=2 offline_train_loop_ARC.py \
  --epochs 100 \
  --loop-core-depth 2 \
  --max-loop-steps 6 \
  --min-loop-steps 2 \
  --batch-size 32 \
  --image-size 64 \
  --patch-size 2 \
  --learning-rate 3e-4 \
  --weight-decay 0 \
  --embed-dim 512 \
  --num-heads 8 \
  --include-rearc \
  --num-colors 12 \
  --data-root "raw_data/ARC-AGI" \
  --train-split "training" \
  --wandb-project "VisionARC" \
  --wandb-run-name "loop_varc_gpu45" \
  --save-path "saves/loop_runs/2025-11-28/checkpoint_final.pt" \
  --best-save-path "saves/loop_runs/2025-11-28/checkpoint_best.pt" \
  --lr-scheduler "cosine" \
  --vis-every 50 \
  --distributed \
  --use-wandb \
  --gate-entropy-weight 5e-4 \
  --loop-penalty-weight 1e-3 \
  --train-dynamic-exit \
  --eval-dynamic-exit
