[ez][AOTI] Fix index offset for Optional Tensor Return (pytorch#155073)

Summary: As title. See added test for more context.

Test Plan:
buck2 run mode/dev-nosan caffe2/test/inductor:test_aot_inductor_custom_ops -- -r test_fn_with_optional_tensor_output_2

Rollback Plan:

Differential Revision: D75900658

Pull Request resolved: pytorch#155073
Approved by: https://github.com/angelayi

Author: yiming0416

test/inductor/custom_ops.cpp

@@ -27,6 +27,22 @@ std::tuple<Tensor, std::optional<Tensor>, std::optional<Tensor>> fn_with_optiona
   return {t3, t4, t5};
 }
 
+std::tuple<Tensor, std::optional<Tensor>, std::optional<Tensor>> fn_with_optional_tensor_output_2_impl(Tensor t1, Tensor t2) {
+  Tensor t3 = t1 + t2;
+  Tensor t4;
+  Tensor t5 = t1 - t2;
+  return {t3, t4, t5};
+}
+
+std::tuple<Tensor, std::optional<Tensor>, std::optional<Tensor>> fn_with_optional_tensor_output_2_meta(Tensor t1, Tensor t2) {
+  Tensor t3 = t1.clone();
+  Tensor t4;
+  Tensor t5 = t1.clone();
+  return {t3, t4, t5};
+}
+
+
+
 Tensor fn_with_all_inputs_impl(
     const Tensor& tensor,
     const c10::List<Tensor>& tensors,
@@ -364,6 +380,7 @@ extern "C" {
 TORCH_LIBRARY(aoti_custom_ops, m) {
   m.def("custom_add(Tensor t1, Tensor t2) -> Tensor");
   m.def("fn_with_optional_tensor_output(Tensor t1, Tensor t2) -> (Tensor, Tensor?, Tensor?)");
+  m.def("fn_with_optional_tensor_output_2(Tensor t1, Tensor t2) -> (Tensor, Tensor?, Tensor?)");
   m.def(
       "fn_with_all_inputs(Tensor tensor, "
       "Tensor[] tensors, "
@@ -410,6 +427,7 @@ TORCH_LIBRARY(aoti_custom_ops, m) {
 TORCH_LIBRARY_IMPL(aoti_custom_ops, CompositeExplicitAutograd, m) {
   m.impl("custom_add", at::custom_add_impl);
   m.impl("fn_with_optional_tensor_output", at::fn_with_optional_tensor_output_impl);
+  m.impl("fn_with_optional_tensor_output_2", at::fn_with_optional_tensor_output_2_impl);
   m.impl("fn_with_all_inputs", at::fn_with_all_inputs_impl);
   m.impl("fn_with_default_input", at::fn_with_default_input_impl);
   m.impl("fn_with_tuple_output", at::fn_with_tuple_output_impl);
@@ -422,6 +440,7 @@ TORCH_LIBRARY_IMPL(aoti_custom_ops, CompositeExplicitAutograd, m) {
 
 TORCH_LIBRARY_IMPL(aoti_custom_ops, Meta, m) {
   m.impl("fn_with_optional_tensor_output", at::fn_with_optional_tensor_output_meta);
+  m.impl("fn_with_optional_tensor_output_2", at::fn_with_optional_tensor_output_2_meta);
   m.impl("fn_with_all_inputs", at::fn_with_all_inputs_meta);
   m.impl("fn_with_default_input", at::fn_with_default_input_meta);
   m.impl("fn_with_tuple_output", at::fn_with_tuple_output_meta);

test/inductor/test_aot_inductor_custom_ops.py

@@ -149,6 +149,18 @@ def forward(self, x, y):
         )
         self.check_model(m, args)
 
+    def test_fn_with_optional_tensor_output_2(self) -> None:
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.ops.aoti_custom_ops.fn_with_optional_tensor_output_2(x, y)
+
+        m = M().to(device=self.device)
+        args = (
+            torch.randn(3, 3, device=self.device),
+            torch.randn(3, 3, device=self.device),
+        )
+        self.check_model(m, args)
+
     def test_custom_op_all_inputs(self) -> None:
         class MyModel(torch.nn.Module):
             # pyre-fixme[3]: Return type must be annotated.

torch/csrc/inductor/aoti_torch/oss_proxy_executor.cpp

@@ -849,7 +849,7 @@ void OSSProxyExecutor::call_function(
           TORCH_CHECK(false, "Expected tensor, got None");
         }
       } else {
-        continue;
+        index++;
       }
     } else {
       TORCH_CHECK(