#!/bin/bash
#SBATCH --job-name=correction
#SBATCH --output=outlook/logs/correction_%j.out
#SBATCH --error=outlook/logs/correction_%j.err
#SBATCH --partition=gpu
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=4
#SBATCH --mem=16G
#SBATCH --time=02:00:00

BATCH_SIZE=${1:-1}

cd /public/home/dxw/Codes/afem

mkdir -p outlook/logs

python -u outlook/src/train_correction.py \
    --data-dir outlook/data_correction \
    --epochs 100 \
    --batch-size "$BATCH_SIZE" \
    --lr 1e-3 \
    --device cuda \
    --checkpoint-out outlook/ckpt/correction.pt