Try remove no_grad in exir capture.

Summary: Remove the no_grad flag in capture to see whether it's still relevant.

Reviewed By: gmagogsfm, guangy10, angelayi

Differential Revision: D48079330

fbshipit-source-id: 38d9158d0a2f1244584a9b165f7020b755ee3d24

Author: ydwu4

exir/capture/_capture.py

@@ -45,7 +45,6 @@
 
 
 @compatibility(is_backward_compatible=False)
-@torch.no_grad()
 def capture(
     f: Callable[..., Any],
     args: Tuple[Value, ...],

exir/tests/test_tracer.py

@@ -306,7 +306,9 @@ def __init__(self):
                 super().__init__()
                 self.register_buffer(
                     "_bin_num_examples",
-                    torch.empty([42]).fill_(0.0),
+                    torch.empty([42]).fill_(
+                        0.0,
+                    ),
                 )
 
             def forward(self, x, y, z):
@@ -327,8 +329,30 @@ def forward(self, x, y, z):
             torch.tensor(3.14),
         )
 
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Found a graph input that requires gradients, and received a mutation.",
+        ):
+            _ = exir.capture(
+                model,
+                example_inputs,
+                exir.CaptureConfig(
+                    pt2_mode=True,
+                    enable_aot=True,
+                ),
+            )
+
+        # Note that model._bin_num_examples is mutated during exir.capture
+        # We need to create a new_model
+        new_model = Module()
+        example_inputs = (
+            torch.randn(4),
+            torch.tensor(0),
+            torch.tensor(3.14),
+        )
+
         ep = exir.capture(
-            model,
+            new_model,
             example_inputs,
             exir.CaptureConfig(
                 pt2_mode=True,
@@ -342,7 +366,7 @@ def forward(self, x, y, z):
             torch.tensor(2.1),
         )
         graph_outputs = ep(*test_inputs)
-        eager_outputs = model(*test_inputs)
+        eager_outputs = new_model(*test_inputs)
         self.assertEqual(len(graph_outputs), 2)
         self.assertEqual(len(eager_outputs), 2)
         self.assertTrue(torch.allclose(graph_outputs[0], eager_outputs[0]))