diff --git a/llmfoundry/callbacks/__init__.py b/llmfoundry/callbacks/__init__.py index 8c86dda2a6..ab028fddcb 100644 --- a/llmfoundry/callbacks/__init__.py +++ b/llmfoundry/callbacks/__init__.py @@ -7,6 +7,7 @@ LRMonitor, MemoryMonitor, MemorySnapshot, + NaNMonitor, OOMObserver, OptimizerMonitor, RuntimeEstimator, @@ -58,6 +59,7 @@ callbacks.register('run_timeout', func=RunTimeoutCallback) callbacks.register('loss_perp_v_len', func=LossPerpVsContextLengthLogger) callbacks.register('env_logger', func=EnvironmentLoggingCallback) +callbacks.register('nan_monitor', func=NaNMonitor) callbacks_with_config.register('async_eval', func=AsyncEval) callbacks_with_config.register('curriculum_learning', func=CurriculumLearning)