#!/bin/bash
# Batch inference over the held-out pairs; 2 GPUs per task (one per expert).
cd "$(dirname "$0")/.."
source env.sh
CKPT=${1:-checkpoints/stage3/final.pt}
export CUDA_VISIBLE_DEVICES=$((SLURM_LOCALID*2)),$((SLURM_LOCALID*2+1))
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
exec .venv/bin/python -m viper.infer_batch \
  --ckpt "$CKPT" --pairs data/pairs_val.jsonl \
  --out_dir results/videos --manifest results/manifest.jsonl \
  --ablate_no_physics --limit ${INFER_LIMIT:-0} \
  --shard $SLURM_LOCALID --num_shards $SLURM_NTASKS