move to device to support pipeline parallel

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -107,7 +107,7 @@ def forward(self, value: Tensor) -> Tensor:
 
         return (
             apply_transform_weight(
-                weight.to(self._precision),
+                weight.to(dtype=self._precision, device=value.device),
                 value.to(self._precision),
                 self.args.location,
                 self.module_type,

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -94,7 +94,7 @@ def __init__(
 
     def forward(self, value: Tensor) -> Parameter:
         return apply_transform_weight(
-            self.weight.to(self._precision),
+            self.weight.to(dtype=self._precision, device=value.device),
             value.to(self._precision),
             self.args.location,
             self.module_type,
@@ -103,7 +103,7 @@ def forward(self, value: Tensor) -> Parameter:
     def right_inverse(self, value: Tensor) -> Tensor:
         inverse = high_precision_invert(self.weight)
         return apply_transform_weight(
-            inverse.to(self._precision),
+            inverse.to(dtype=self._precision, device=value.device),
             value.to(self._precision),
             self.args.location,
             self.module_type,