[DO NOT MERGE] [AOT Autograd] Trace the backward pass assuming contiguous tensors

Author: anijain2305

functorch/_src/aot_autograd.py

@@ -135,7 +135,7 @@ def forward(ctx, *flat_tensor_args):
                 with torch.set_grad_enabled(grad_state):
                     out = flat_fn(*flat_tensor_args)
                 out = pytree.tree_map(
-                    lambda x: x.detach() if isinstance(x, Tensor) else x, out
+                    lambda x: x.detach().contiguous() if isinstance(x, Tensor) else x, out
                 )
 
                 if isinstance(out, (list, tuple)):
@@ -164,9 +164,8 @@ def forward(ctx, *flat_tensor_args):
 
         @staticmethod
         def backward(ctx, *flat_args):
-            # hmm... this doesn't feel right. todo
-            # contiguous_args = [t.contiguous() for t in flat_args]
-            contiguous_args = [t for t in flat_args]
+            contiguous_args = [t.contiguous() for t in flat_args]
+            # contiguous_args = [t for t in flat_args]
             out = normalize_as_list(compiled_bw(*ctx.saved_tensors, *contiguous_args))
             return tuple(out)