add pre-emption callback to esm2 train

NVIDIA · Nov 14, 2024 · d9177e0 · d9177e0
1 parent 78d9a2e
commit d9177e0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py b/sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py
@@ -190,6 +190,7 @@ def main(
         PerplexityLoggingCallback(log_train=False, log_val=True),
         RichModelSummary(max_depth=4),
         LearningRateMonitor(),
+        nl_callbacks.PreemptionCallback(),
     ]
     if nsys_profiling:
         if nsys_end_step is None: