Fix CI

pytorch · Nov 20, 2024 · 40cabd1 · 40cabd1
1 parent bac5d4a
commit 40cabd1
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchao/quantization/qat/linear.py b/torchao/quantization/qat/linear.py
@@ -365,6 +365,7 @@ def _convert_qat_linear_4w(self, module: torch.nn.Module):
                     inner_k_tiles=inner_k_tiles,
                     precision=child.weight.dtype,
                     scales_precision=config.scale_precision,
+                    device=next(child.parameters()).device,
                 )
                 setattr(module, name, quantized_linear)