Update ssd.py (#33)

Ensure the variables are on the same device.
OpenPPL · Mar 17, 2022 · f435f35 · f435f35
1 parent 97a5555
commit f435f35
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ppq/quantization/optim/ssd.py b/ppq/quantization/optim/ssd.py
@@ -230,7 +230,7 @@ def one_step_equalization(
             kernel_scale = kernel_scale / next_kernel_scale
             act_scale = act_scale / next_kernel_scale
             min_scale = torch.min(kernel_scale, act_scale)
-            min_scale = torch.min(min_scale, torch.tensor(default_min_scale, dtype=torch.float32))
+            min_scale = torch.min(min_scale, torch.tensor(default_min_scale, dtype=torch.float32, device=min_scale.device))
             min_scale /= min_scale.min()
             min_scale = torch.clip(min_scale, 1.0, max_scale)
         else:
@@ -482,4 +482,4 @@ def optimize(
                     logger.info('SSD and DFQ Deactivated')
                 self.initiate_pair_state(pair)
 
-
+