fix staging callback blocking bug

tensorpack/graph_builder/input_source.py

@@ -411,6 +411,7 @@ class StagingInputWrapper(FeedfreeInput):
 
         def _before_train(self):
             # pre-fill the staging area
+            logger.info("Pre-fillilng staging area ...")
             for k in range(self.nr_stage):
                 self.stage_op.run()
 

tensorpack/train/base.py

@@ -114,8 +114,7 @@ class Trainer(object):
         """
         input_source.setup(self.model.get_inputs_desc())
         cbs = input_source.get_callbacks()
-        for cb in cbs:
-            self.register_callback(cb)
+        self.config.callbacks.extend(cbs)
 
     def setup(self):
         """