Enhance flexibility by passing callbacks as method argument

Signed-off-by: Michal Futrega <[email protected]>
NVIDIA · Dec 13, 2023 · 5e24036 · 5e24036
1 parent af8daed
commit 5e24036
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/collections/nlp/parts/megatron_trainer_builder.py b/nemo/collections/nlp/parts/megatron_trainer_builder.py
@@ -95,10 +95,10 @@ def _plugins(self) -> list:
 
         return plugins
 
-    def create_trainer(self) -> Trainer:
+    def create_trainer(self, callbacks=(CustomProgressBar(),)) -> Trainer:
         strategy = self._training_strategy()
         plugins = self._plugins()
-        return Trainer(plugins=plugins, strategy=strategy, **self.cfg.trainer, callbacks=[CustomProgressBar()])
+        return Trainer(plugins=plugins, strategy=strategy, **self.cfg.trainer, callbacks=callbacks)
 
 
 class MegatronBertTrainerBuilder(MegatronTrainerBuilder):