add logging when finished

Signed-off-by: Gerald Shen <[email protected]>
NVIDIA · Dec 8, 2023 · 07e3990 · 07e3990
1 parent 39576f6
commit 07e3990
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/nemo_aligner/algorithms/supervised.py b/nemo_aligner/algorithms/supervised.py
@@ -18,6 +18,8 @@
 import torch
 from omegaconf.dictconfig import DictConfig
 from tqdm import tqdm
+from nemo.utils import logging
+
 
 from nemo_aligner.utils.distributed import SyncTimer
 from nemo_aligner.utils.train_utils import clip_gradients
@@ -198,6 +200,7 @@ def fit(self):
                     self.save(metrics, is_train_end=is_train_end)
 
                 if run_time_exceeded:
+                    logging.info(f"Time limit given by run_timer={self.run_timer} reached. Stopping run")
                     return
 
                 metrics.clear()