documentation (#361)

* documentation

* foc

Author: xadupre

_unittests/ut_export/test_cf_simple_loop_for.py

@@ -104,7 +104,9 @@ def body(i: torch.Tensor, x: torch.Tensor) -> Tuple[torch.Tensor]:
     def test_simple_loop_for_2(self):
         class Model(torch.nn.Module):
             def forward(self, n_iter, x):
-                def body(i: torch.Tensor, x: torch.Tensor) -> Tuple[torch.Tensor]:
+                def body(
+                    i: torch.Tensor, x: torch.Tensor
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
                     return (x[: i.item() + 1].unsqueeze(1), x[i.item() + 1 :].unsqueeze(1))
 
                 return simple_loop_for(n_iter, body, (x,))
@@ -172,8 +174,13 @@ def body(i: torch.Tensor, x: torch.Tensor) -> Tuple[torch.Tensor]:
     def test_simple_loop_for_2_concatenation_dims(self):
         class Model(torch.nn.Module):
             def forward(self, n_iter, x):
-                def body(i: torch.Tensor, x: torch.Tensor) -> Tuple[torch.Tensor]:
-                    return (x[: i.item() + 1].unsqueeze(1), x[i.item() + 1 :].unsqueeze(0))
+                def body(
+                    i: torch.Tensor, x: torch.Tensor
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+                    return (
+                        x[: i.item() + 1].unsqueeze(1),
+                        x[i.item() + 1 :].unsqueeze(0),
+                    )
 
                 return simple_loop_for(n_iter, body, (x,), (0, 1))
 

onnx_diagnostic/export/cf_simple_loop_for.py

@@ -321,7 +321,7 @@ def body(i, x):
 
         class Model(torch.nn.Module):
             def forward(self, n_iter, x):
-                def body(i: torch.Tensor, x: torch.Tensor) -> Tuple[torch.Tensor]:
+                def body(i, x):
                     return (x[: i.item() + 1].unsqueeze(1), x[i.item() + 1 :].unsqueeze(0))
 
                 return simple_loop_for(n_iter, body, (x,), (0, 1))
@@ -346,6 +346,7 @@ def body(i: torch.Tensor, x: torch.Tensor) -> Tuple[torch.Tensor]:
         ),
     )
     torch._check(
-        isinstance(res, tuple), f"Output of the loop should be a tuple not {type(res)}."
+        isinstance(res, tuple),
+        lambda: f"Output of the loop should be a tuple not {type(res)}.",
     )
     return res[0] if len(res) == 1 else res