Enhance flexibility by passing callbacks as method argument (NVIDIA#8015

) * Enhance flexibility by passing callbacks as method argument Signed-off-by: Michal Futrega <michal.futrega@gmail.com> * Set callbacks default to None Signed-off-by: Michal Futrega <michal.futrega@gmail.com> --------- Signed-off-by: Michal Futrega <michal.futrega@gmail.com>
rohitrango · Jan 9, 2024 · 0de93a7 · 0de93a7
1 parent 4b7f7ea
commit 0de93a7
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/nemo/collections/nlp/parts/megatron_trainer_builder.py b/nemo/collections/nlp/parts/megatron_trainer_builder.py
@@ -118,10 +118,12 @@ def _plugins(self) -> list:
 
         return plugins
 
-    def create_trainer(self) -> Trainer:
+    def create_trainer(self, callbacks=None) -> Trainer:
         strategy = self._training_strategy()
         plugins = self._plugins()
-        return Trainer(plugins=plugins, strategy=strategy, **self.cfg.trainer, callbacks=[CustomProgressBar()])
+        if callbacks is None:
+            callbacks = [CustomProgressBar()]
+        return Trainer(plugins=plugins, strategy=strategy, **self.cfg.trainer, callbacks=callbacks)
 
 
 class MegatronBertTrainerBuilder(MegatronTrainerBuilder):