YOLOV8 Fix modify arch for quantization call (#1686)

neuralmagic · Jul 25, 2023 · ab16210 · ab16210
1 parent 09c9c96
commit ab16210
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/sparseml/yolov8/trainers.py b/src/sparseml/yolov8/trainers.py
@@ -233,7 +233,7 @@ def _build_managers(self, ckpt: Optional[dict]):
             )
             self.checkpoint_manager = ScheduledModifierManager.from_yaml(ckpt["recipe"])
             if self.checkpoint_manager.quantization_modifiers:
-                self._modify_arch_for_quantization()
+                _modify_arch_for_quantization(self.model)
             self.checkpoint_manager.apply_structure(self.model, epoch=float("inf"))
 
         else:
@@ -245,7 +245,7 @@ def _build_managers(self, ckpt: Optional[dict]):
                 f"at epoch {ckpt['epoch']}"
             )
             if self.manager.quantization_modifiers:
-                self._modify_arch_for_quantization()
+                _modify_arch_for_quantization(self.model)
             self.manager.apply_structure(self.model, epoch=ckpt["epoch"])
 
     def resume_training(self, ckpt):