Forward fix NluStellaCapExportTest

Summary:
Forward fix NluStellaCapExportTest. Currently, to_backend works with graph module and doesn't work well with ExportedProgram. Need to return the graph_module contained in ExportedProgram to fix the test.

There is an assertion error raised in call_delegate_autograd: it finds some inputs require grad. In this diff, we delay the error message to when .backward() is called (probably will never be called since we're in inference mode, which is good since we only care about inference for call_delegate).

I checked how nlu models generate input and find no inputs requires grad. I suspect that there are some transformations doesn't respect this requires_grad=False property. This deserves a separate investigation.

Reviewed By: angelayi

Differential Revision: D46995565

fbshipit-source-id: 3d735843b8db0fe2c003a74cb9895c68b6cdac51

Author: ydwu4

exir/delegate.py

@@ -152,14 +152,34 @@ def call_delegate_cpu(lowered_module, *args):
 def call_delegate_autograd(lowered_module, *args):
     # TODO: support autograd
     flat_operands, _ = tree_flatten([lowered_module, *args])
-    assert all(
-        [not f.requires_grad for f in flat_operands if isinstance(f, torch.Tensor)]
+    requires_grad = any(
+        [f.requires_grad for f in flat_operands if isinstance(f, torch.Tensor)]
     )
 
-    _ = torch._C.ExcludeDispatchKeyGuard(
+    with torch._C._ExcludeDispatchKeyGuard(
         torch._C.DispatchKeySet(torch._C.DispatchKey.AutogradCPU)
-    )
-    return executorch_call_delegate(lowered_module, *args)
+    ):
+        res = executorch_call_delegate(lowered_module, *args)
+
+        if requires_grad:
+            err_fn = torch._C._functions.DelayedError(
+                b"NYI: call_delegate doesn't support autograd",
+                1,
+            )
+            # Create aliases of the output that has requires_grad=True. We need
+            # at least one of the inputs to err_fn to require grad so that the
+            # output will have a grad_fn.
+
+            # pyre-ignore
+            def fake_requires_grad(var):
+                if var is not None:
+                    var = var.detach()
+                    var.requires_grad = True
+                return err_fn(var)
+
+            return pytree.tree_map(fake_requires_grad, res)
+
+        return res
 
 
 @executorch_call_delegate.py_impl(ProxyTorchDispatchMode)