#!/bin/bash

#SBATCH --job-name=tf-checkpoint
#SBATCH --partition=main
#SBATCH --gpus=1
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --mem=20G
#SBATCH --time=00:45:00
#SBATCH --output=slurm.%N.%j.out
#SBATCH --error=slurm.%N.%j.err

# Training variables
script=$PWD/example-04-checkpoint.py
ckpt_dir=$PWD/models_tf

# Hyperparameters
batch_size=32
epochs=10
lr=0.01
momentum=0.5

apptainer exec docker://tensorflow/tensorflow:latest \
    python "$script" \
    --ckpt-dir "$ckpt_dir" \
    --batch-size "$batch_size" \
    --epochs "$epochs" \
    --lr "$lr" \
    --momentum "$momentum" \
    --resume-training