specify smoothquant device (#1830)

neuralmagic · Nov 16, 2023 · 7b6ac03 · 7b6ac03
1 parent 5e76668
commit 7b6ac03
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