add suggested fix from david

stanford-crfm · Oct 18, 2024 · 7c195ba · 7c195ba
1 parent 37e77fb
commit 7c195ba
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/levanter/callbacks.py b/src/levanter/callbacks.py
@@ -39,6 +39,8 @@ def log_epoch(step_info: StepInfo):
         nonlocal total_tokens
         if total_tokens is None:
             if not total_tokens_future.done():
+                if step_info.step % 1000 == 0:
+                    logger.info("Dataset not finished. Can't compute epochs.")
                 return  # We don't have the total tokens yet, so we can't calculate epoch
             total_tokens = total_tokens_future.result()