FakeTensor should be wrapped as TensorVariable (#931)

Author: yanboliang

test/test_misc.py

@@ -2060,6 +2060,21 @@ def fn(x, m):
         res = opt_fn(x, m)
         self.assertEqual(ref, res)
 
+    def test_tensor_dot_grad_no_graph_break(self):
+        def fn(a, b):
+            y = 3 * a**3 - b**2
+            y.backward(gradient=torch.tensor([1.0, 1.0]))
+            b.grad.zero_()
+            return a.grad, b.grad
+
+        a = torch.tensor([2.0, 3.0], requires_grad=True)
+        b = torch.tensor([6.0, 4.0], requires_grad=True)
+        cnts = torchdynamo.testing.CompileCounter()
+        with torchdynamo.optimize(cnts):
+            _, b_grad = fn(a, b)
+        self.assertTrue(same(b_grad, torch.tensor([0.0, 0.0])))
+        self.assertEqual(cnts.frame_count, 2)
+
     def test_change_backends(self):
         @torchdynamo.optimize("eager", nopython=True)
         def fn1():

torchdynamo/utils.py

@@ -188,9 +188,14 @@ def is_numpy_float_type(value):
 
 def istensor(obj):
     """Check of obj is a tensor"""
-    return istype(
-        obj, (torch.Tensor, torch.nn.Parameter, *config.traceable_tensor_subclasses)
+    tensor_list = (
+        torch.Tensor,
+        torch.nn.Parameter,
+        *config.traceable_tensor_subclasses,
     )
+    if fake_tensors_available:
+        tensor_list = tensor_list + (torch._subclasses.FakeTensor,)
+    return istype(obj, tensor_list)
 
 
 def is_lazy_module(mod):