From 354a6419b6ab9f6c787d84131090e03245f240ba Mon Sep 17 00:00:00 2001
From: Henry Addison <henryaddison@users.noreply.github.com>
Date: Mon, 18 Mar 2024 15:42:47 +0000
Subject: [PATCH] use checkpoint rather than epoch for determining checkpoint
 for sampling

just cleaner and more consistent
---
 bin/predict.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/bin/predict.py b/bin/predict.py
index adcc09d71..c6eb23da1 100644
--- a/bin/predict.py
+++ b/bin/predict.py
@@ -239,7 +239,7 @@ def main(
     workdir: Path,
     dataset: str = typer.Option(...),
     split: str = "val",
-    epoch: int = typer.Option(...),
+    checkpoint: str = typer.Option(...),
     batch_size: int = None,
     num_samples: int = 3,
     input_transform_dataset: str = None,
@@ -260,7 +260,7 @@ def main(
 
     output_dirpath = samples_path(
         workdir=workdir,
-        checkpoint=f"epoch-{epoch}",
+        checkpoint=checkpoint,
         dataset=dataset,
         input_xfm=f"{config.data.input_transform_dataset}-{config.data.input_transform_key}",
         split=split,
@@ -290,7 +290,7 @@ def main(
         shuffle=False,
     )
 
-    ckpt_filename = os.path.join(workdir, "checkpoints", f"epoch_{epoch}.pth")
+    ckpt_filename = os.path.join(workdir, "checkpoints", f"{checkpoint}.pth")
     logger.info(f"Loading model from {ckpt_filename}")
     state, sampling_fn = load_model(config, ckpt_filename)