Assert that there are no None arguments to backwards

ezyang · ezyang · commit c8c06a89511e · 2023-03-08T08:44:52.000-05:00
This assert would have caught #96219 Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]
diff --git a/torch/_functorch/aot_autograd.py b/torch/_functorch/aot_autograd.py
@@ -2328,6 +2328,7 @@ def backward(ctx, *flat_args):
 
             def call_compiled_backward():
                 if CompiledFunction.compiled_bw is None:
+                    assert all(a is not None for a in all_args)
                     if aot_config.dynamic_shapes:
                         all_args_list = list(all_args)
                         CompiledFunction.compiled_bw = create_aot_dispatcher_function(
