specify smoothquant device (#1830)

neuralmagic · Nov 14, 2023 · 14ec264 · 14ec264
1 parent d4a9ce9
commit 14ec264
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/sparseml/modifiers/smoothquant/pytorch.py b/src/sparseml/modifiers/smoothquant/pytorch.py
@@ -127,6 +127,7 @@ def _calibrate(self, model: ModifiableModelPyTorch, calibration_dataloader: List
             calibration_dataloader,
             self.num_calibration_steps,
             self.calibration_function,
+            self.device_,
         )
 
         # remove the hooks now that we are done calibrating