fix spurious aot autograd warning (#95521) (#95614)

bdhirsh · web-flow · commit 66d23dbad7c8 · 2023-02-27T14:31:02.000-05:00
The _make_boxed logic probably needs a cleanup, but this fixes a spurious warning that we should get in before the release. Confirmed that this used to emit a warning and no longer does: ``` import torch lin = torch.nn.Linear(100, 10) def f(x): return lin(x) opt_f = torch.compile(f) opt_f(torch.randn(10, 100, requires_grad=False)) ``` Pull Request resolved: #95521 Approved by: https://github.com/ngimel
diff --git a/torch/_functorch/aot_autograd.py b/torch/_functorch/aot_autograd.py
@@ -1870,6 +1870,9 @@ def create_runtime_wrapper(
     trace_joint: bool,
     keep_input_mutations: bool,
 ):
+    if not hasattr(compiled_fn, "_boxed_call"):
+        compiled_fn = make_boxed_func(compiled_fn)
+
     def runtime_wrapper(*args):
         # Step 2: remove aliased inputs that are mutated, replace with synthetic bases
         # Only happens if our graph mutates an input that aliases another input.
