bugfix: set_epoch after reume

PaddlePaddle · Mar 6, 2023 · e6b4f1e · e6b4f1e
1 parent b98d88e
commit e6b4f1e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ppcls/engine/engine.py b/ppcls/engine/engine.py
@@ -385,6 +385,9 @@ def train(self):
                                      ema_module)
             if metric_info is not None:
                 best_metric.update(metric_info)
+            if hasattr(self.train_dataloader.batch_sampler, "set_epoch"):
+                self.train_dataloader.batch_sampler.set_epoch(best_metric[
+                    "epoch"])
 
         for epoch_id in range(best_metric["epoch"] + 1,
                               self.config["Global"]["epochs"] + 1):