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

cd /public/home/dxw/Codes/afem

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