Fix tensor.data access under inference_mode and compile (#134878)

zou3519 · pytorchmergebot · commit d7b57c4d63ed · 2024-09-04T17:55:41.000Z
Fixes #134798 In the regular Tensor case, when you call Tensor.data, there's a check for if inference mode is active. If it is active, then we don't set the version counter. We replicate this check for Tensor Subclasses (the bug was we were trying to set the version counter on a FakeTensor in inference_mode). Test Plan: - new test Pull Request resolved: #134878 Approved by: https://github.com/bdhirsh
diff --git a/c10/core/TensorImpl.cpp b/c10/core/TensorImpl.cpp
@@ -509,7 +509,9 @@ c10::intrusive_ptr<TensorImpl> TensorImpl::shallow_copy_and_detach_core(
     r = (pyobj_slot_.load_pyobj_interpreter())->detach(this);
   }
   if (r) {
-    r->set_version_counter(std::forward<VariableVersion>(version_counter));
+    if (!r->is_inference()) {
+      r->set_version_counter(std::forward<VariableVersion>(version_counter));
+    }
     r->set_allow_tensor_metadata_change(allow_tensor_metadata_change);
     return r;
   }
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
@@ -2000,6 +2000,17 @@ def fn(cfg, x, y):
         self.assertEqual(cnts.frame_count, 1)
         self.assertEqual(cnts.op_count, 2)
 
+    def test_data_access_in_inference_mode(self):
+        @torch.compile(fullgraph=True)
+        def f(x):
+            y = x.data
+            return y
+
+        with torch.inference_mode():
+            x = torch.randn(3)
+            y = f(x)
+        self.assertEqual(y, x)
+
     def test_dataclass_fields(self):
         @dataclasses.dataclass
         class MyDataClass:

Original file line number	Diff line number	Diff line change
`@@ -509,7 +509,9 @@ c10::intrusive_ptr<TensorImpl> TensorImpl::shallow_copy_and_detach_core(`
`509`	`509`	`r = (pyobj_slot_.load_pyobj_interpreter())->detach(this);`
`510`	`510`	`}`
`511`	`511`	`if (r) {`
`512`		`- r->set_version_counter(std::forward<VariableVersion>(version_counter));`
	`512`	`+ if (!r->is_inference()) {`
	`513`	`+ r->set_version_counter(std::forward<VariableVersion>(version_counter));`
	`514`	`+ }`
`513`	`515`	`r->set_allow_tensor_metadata_change(allow_tensor_metadata_change);`
`514`	`516`	`return r;`
`515`	`517`	`}`