diff --git a/src/sparseml/yolov8/trainers.py b/src/sparseml/yolov8/trainers.py
index 46398bcfb94..e8c4c830996 100644
--- 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):