add trainer.sanity_check

NVIDIA · Oct 15, 2024 · c8a7d1a · c8a7d1a
1 parent eaabb31
commit c8a7d1a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py b/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py
@@ -360,6 +360,9 @@ def on_megatron_reduce_microbatches_end(
             - forward_out: dict of tensors with the following keys:
                 - token_logits: [b s vocab]
         """
+        if step.trainer.sanity_checking:  # skip sanity check
+            return
+
         if step.trainer.training and not self.log_train:
             return