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

ebsmothers · web-flow · commit 0ed30902e387 · 2024-08-23T01:51:25.000-04:00
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)
 

Original file line number	Diff line number	Diff line change
`@@ -339,7 +339,7 @@ def copy_(func, args, *kwargs):`
`339`	`339`	`# Convert Non NF4Tensor into NF4 for copy in`
`340`	`340`	`if not isinstance(copy_in, NF4Tensor):`
`341`	`341`	`copy_in_nf4 = NF4Tensor.from_tensor(`
`342`		`- copy_in, original.block_size, original.scaler_block_size`
	`342`	`+ copy_in.to(original.device), original.block_size, original.scaler_block_size`
`343`	`343`	`)`
`344`	`344`	`return original.copy_(copy_in_nf4)`
`345`	`345`