Update on "Add missing in-place on view check to custom autograd.Function"

soulitzer · soulitzer · commit d5c894593c1a · 2025-05-09T12:11:10.000-04:00
Fixes #152773 [ghstack-poisoned]
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -8673,6 +8673,26 @@ def backward(ctx, gO):
         ):
             Func.apply(b)
 
+    def test_custom_function_inplace_on_view_of_leaf(self):
+        class Func(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, inp):
+                inp.add_(1)
+                ctx.mark_dirty(inp)
+                return inp
+
+            @staticmethod
+            def backward(ctx, gO):
+                pass
+
+        a = torch.tensor([1.0, 2.0], requires_grad=True)
+        b = a.view_as(a)
+
+        with self.assertRaisesRegex(
+            RuntimeError, "a view of a leaf Variable that requires grad"
+        ):
+            Func.apply(b)
+
     def test_named_tensor_for_complex_views(self):
         names = ["batch", "height", "width", "complex"]
         z = torch.ones((2, 1, 2, 2), requires_grad=True)
diff --git a/torch/csrc/autograd/custom_function.cpp b/torch/csrc/autograd/custom_function.cpp
@@ -338,7 +338,7 @@ static optional_variable_list _process_backward_mode_ad(
         auto& grad_acc = dynamic_cast<AccumulateGrad&>(*grad_acc_fn);
         grad_acc.variable.reset();
       }
-      // This repeats checks like mutation of leaf variables already done above
+      // This repeats the mutation of leaf variables check already done above
       check_inplace(var, true);
       impl::rebase_history(var, {cdata, output_nr});
     } else if (is_input) {
