Merge pull request #692 from mv1388/callback_handler_mp_filter_fix

Fix mp filtering the callbacks - add train loop callbacks list filtering
mv1388 · Jul 12, 2022 · d4e265a · d4e265a
2 parents 3a562e7 + aa1381c
commit d4e265a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/aitoolbox/torchtrain/train_loop/components/callback_handler.py b/aitoolbox/torchtrain/train_loop/components/callback_handler.py
@@ -159,7 +159,8 @@ def split_on_execution_position(self, callbacks, register_train_loop=False):
                 cbs_at_position.sort(key=lambda cb: cb.execution_order)
 
     def mp_filter_callbacks(self):
-        self.mp_filter_callbacks()
+        self.train_loop_obj.callbacks = self._mp_filter_cb_list(self.train_loop_obj.callbacks)
+
         self.cbs_on_epoch_begin = self._mp_filter_cb_list(self.cbs_on_epoch_begin)
         self.cbs_on_epoch_end = self._mp_filter_cb_list(self.cbs_on_epoch_end)
         self.cbs_on_train_begin = self._mp_filter_cb_list(self.cbs_on_train_begin)