Add load checkpoint callback (#1570)

mosaicml · Oct 28, 2024 · 3eb5cea · 3eb5cea
1 parent a3ddced
commit 3eb5cea
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/llmfoundry/callbacks/__init__.py b/llmfoundry/callbacks/__init__.py
@@ -4,6 +4,7 @@
 from composer.callbacks import (
     EarlyStopper,
     Generate,
+    LoadCheckpoint,
     LRMonitor,
     MemoryMonitor,
     MemorySnapshot,
@@ -63,6 +64,7 @@
 callbacks.register('env_logging', func=EnvironmentLoggingCallback)
 callbacks.register('nan_monitor', func=NaNMonitor)
 callbacks.register('kill_loss_spike', func=KillLossSpike)
+callbacks.register('load_checkpoint', func=LoadCheckpoint)
 
 callbacks_with_config.register('async_eval', func=AsyncEval)
 callbacks_with_config.register('curriculum_learning', func=CurriculumLearning)