feat: add on_epoch_end_cb to call at end of each epoch

torchensemble/soft_gradient_boosting.py

@@ -209,6 +209,7 @@ def fit(
         test_loader=None,
         save_model=True,
         save_dir=None,
+        on_epoch_end_cb=None
     ):
 
         # Instantiate base estimators and set attributes
@@ -297,6 +298,9 @@ def fit(
                 else:
                     scheduler.step()
 
+            # Call on epoch end
+            if on_epoch_end_cb:
+                on_epoch_end_cb(epoch)
         if save_model and not test_loader:
             io.save(self, save_dir, self.logger)
 
@@ -392,6 +396,7 @@ def fit(
         test_loader=None,
         save_model=True,
         save_dir=None,
+        on_epoch_end_cb=None
     ):
         super().fit(
             train_loader=train_loader,
@@ -401,6 +406,7 @@ def fit(
             test_loader=test_loader,
             save_model=save_model,
             save_dir=save_dir,
+            on_epoch_end_cb=on_epoch_end_cb,
         )
 
     @torchensemble_model_doc(