include kwargs

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

test/dtypes/test_nf4.py

@@ -214,8 +214,8 @@ def test_smoketest_linear(self, dtype: torch.dtype):
     def test_smoketest_linear_compile(self, dtype: torch.dtype):
         if torch.cuda.is_available() and torch.cuda.get_device_capability() < (8, 0) and dtype == torch.bfloat16:
             self.skipTest("test requires SM capability of at least (8, 0).")
-        if version.parse(torch.__version__) <= version.parse("2.2.2"):
-            self.skipTest("test requires 2.3.0+ for tracing NF4Tensor")
+        if version.parse(torch.__version__) < version.parse("2.3.0"):
+            self.skipTest("test requires 2.3.0 and above for tracing NF4Tensor")
         a = torch.randn(32, 32, dtype=dtype, device='cuda')
         a_nf4 = torchao.dtypes.to_nf4(a, 16, 2)
         inp = torch.randn(2, 32, 32, dtype=a.dtype, device=a.device)

torchao/dtypes/nf4tensor.py

@@ -610,4 +610,12 @@ def decorator(func):
 
 @implements_torch_function(torch.Tensor.to)
 def function_to_dtype(*args, **kwargs):
-    return args[0].get_original_weight().to(args[1])
+    if isinstance(args[0], NF4Tensor) and isinstance(args[1], torch.dtype):
+        # Tensor.to(dtype, non_blocking, copy, memory_format)
+        return args[0].get_original_weight().to(*args[1:], **kwargs)
+    else:
+        # Tensor.to(device, dtype, non_blocking, copy, memory_format)
+        # Tensor.to(other, non_blocking, copy)
+        raise NotImplementedError(
+            f"NF4Tensor.to({args[1:]}, {kwargs}) is not supported"
+        )