check for CKPT when target metric is reached before saving

2026-04-29 03:00:14 -04:00 · 2025-03-02 00:41:08 -08:00
parent 3ac4ae5870
commit 27ec792c19
1 changed files with 1 additions and 1 deletions
--- a/examples/mlperf/model_train.py
+++ b/examples/mlperf/model_train.py
@@ -595,7 +595,7 @@ def train_retinanet():

      if val_metric >= target_metric:
        print(colored(f"target metric reached: {val_metric:.2f}/{target_metric:.2f}"))
-        safe_save(get_state_dict(model), fn)
+        if getenv("CKPT", 1): safe_save(get_state_dict(model), fn)
        break

 def train_unet3d():