Fix CI

yanbing-j · yanbing-j · commit 1f6d59d04f21 · 2024-11-19T07:02:09.000Z
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)
 

Original file line number	Diff line number	Diff line change
`@@ -365,6 +365,7 @@ def _convert_qat_linear_4w(self, module: torch.nn.Module):`
`365`	`365`	`inner_k_tiles=inner_k_tiles,`
`366`	`366`	`precision=child.weight.dtype,`
`367`	`367`	`scales_precision=config.scale_precision,`
	`368`	`+ device=next(child.parameters()).device,`
`368`	`369`	`)`
`369`	`370`	`setattr(module, name, quantized_linear)`
`370`	`371`