Move non-NF4 tensor to device prior to quantization on copy (#737)

pytorch · Aug 23, 2024 · 0ed3090 · 0ed3090
1 parent 68e4643
commit 0ed3090
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchao/dtypes/nf4tensor.py b/torchao/dtypes/nf4tensor.py
@@ -339,7 +339,7 @@ def copy_(func, *args, **kwargs):
     # Convert Non NF4Tensor into NF4 for copy in
     if not isinstance(copy_in, NF4Tensor):
         copy_in_nf4 = NF4Tensor.from_tensor(
-            copy_in, original.block_size, original.scaler_block_size
+            copy_in.to(original.device), original.block_size, original.scaler_block_size
         )
         return original.copy_(copy_in_nf4)