#!/usr/bin/env bash
set -euo pipefail
source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
extra_args=()
case "${1:-}" in
  "") ;;
  --smoke) extra_args=(--max_steps 1) ;;
  *) echo 'Usage: scripts/train.sh [--smoke]' >&2; exit 2 ;;
esac
[[ $# -le 1 ]] || { echo 'Too many arguments' >&2; exit 2; }
export CUDA_VISIBLE_DEVICES="${GPUS:-${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}}"
LOGDIR="$(realpath -m "${LOGDIR:-$REPRO_ROOT/runs/three-encoder}")"
mkdir -p "$LOGDIR"
cp "$CONFIG" "$LOGDIR/config.yaml"
cd "$REPRO_ROOT/code"
"$PYTHON_BIN" -m torch.distributed.run --standalone --nnodes=1 --nproc-per-node=8 train.py \
  --config_path "$LOGDIR/config.yaml" --logdir "$LOGDIR" \
  --disable-wandb --no_visualize "${extra_args[@]}" 2>&1 | tee -a "$LOGDIR/train.log"