fix(sGBM): missing functionality of param save_model

TorchEnsemble-Community · Aug 25, 2021 · 01a9e4c · 01a9e4c
1 parent 02703c2
commit 01a9e4c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchensemble/soft_gradient_boosting.py b/torchensemble/soft_gradient_boosting.py
@@ -278,7 +278,7 @@ def fit(
             # Validation
             if test_loader:
                 flag = self._evaluate_during_fit(test_loader, epoch)
-                if flag:
+                if save_model and flag:
                     io.save(self, save_dir, self.logger)
 
             # Update the scheduler