that kwargs (#1020)

mosaicml · Mar 10, 2024 · 36ab1ba · 36ab1ba
1 parent 2b17497
commit 36ab1ba
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llmfoundry/utils/builders.py b/llmfoundry/utils/builders.py
@@ -165,7 +165,7 @@ def build_callback(
     if name == 'lr_monitor':
         return LRMonitor()
     elif name == 'memory_monitor':
-        return MemoryMonitor()
+        return MemoryMonitor(**kwargs)
     elif name == 'oom_observer':
         return OOMObserver(**kwargs)
     elif name == 'memory_snapshot':