support sparse evaluation (PaddlePaddle#5122)

an1018 · Jan 19, 2022 · 3233ef3 · 3233ef3
1 parent a794cac
commit 3233ef3
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/ppdet/engine/trainer.py b/ppdet/engine/trainer.py
@@ -130,9 +130,10 @@ def __init__(self, cfg, mode='train'):
             self.lr = create('LearningRate')(steps_per_epoch)
             self.optimizer = create('OptimizerBuilder')(self.lr, self.model)
 
-        if self.cfg.get('unstructured_prune'):
-            self.pruner = create('UnstructuredPruner')(self.model,
-                                                       steps_per_epoch)
+            # Unstructured pruner is only enabled in the train mode.
+            if self.cfg.get('unstructured_prune'):
+                self.pruner = create('UnstructuredPruner')(self.model,
+                                                           steps_per_epoch)
 
         self._nranks = dist.get_world_size()
         self._local_rank = dist.get_rank()