#!/usr/bin/env bash
# 启动 HRM Sudoku 1k 训练 (HRM 官方推荐配置)
# 单 GPU 约 10h on RTX 4070; A6000 应该更快
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate rrm

cd "$REPO_ROOT/hrm"
OMP_NUM_THREADS=${OMP_NUM_THREADS:-8} \
WANDB_MODE=${WANDB_MODE:-online} \
python pretrain.py \
  data_path="$REPO_ROOT/data/sudoku-extreme-1k-aug-1000" \
  epochs=20000 eval_interval=2000 global_batch_size=384 \
  lr=7e-5 puzzle_emb_lr=7e-5 weight_decay=1.0 puzzle_emb_weight_decay=1.0 "$@"