diff --git a/llmfoundry/callbacks/__init__.py b/llmfoundry/callbacks/__init__.py index 4712de5d5e..66255d636c 100644 --- a/llmfoundry/callbacks/__init__.py +++ b/llmfoundry/callbacks/__init__.py @@ -11,6 +11,7 @@ OptimizerMonitor, RuntimeEstimator, SpeedMonitor, + SystemMetricsMonitor, ) from llmfoundry.callbacks.async_eval_callback import AsyncEval @@ -35,6 +36,7 @@ from llmfoundry.callbacks.scheduled_gc_callback import ScheduledGarbageCollector from llmfoundry.registry import callbacks, callbacks_with_config +callbacks.register('system_metrics_monitor', func=SystemMetricsMonitor) callbacks.register('lr_monitor', func=LRMonitor) callbacks.register('memory_monitor', func=MemoryMonitor) callbacks.register('memory_snapshot', func=MemorySnapshot) @@ -71,4 +73,4 @@ 'AsyncEval', 'CurriculumLearning', 'LossPerpVsContextLengthLogger', -] +] \ No newline at end of file