fix

Author: xadupre

onnx_diagnostic/export/control_flow.py

@@ -10,7 +10,7 @@
 from .api import to_onnx
 
 _TEST_EXPORT = False
-_REGISTERED_SCHEMA = {}
+_REGISTERED_SCHEMA = {}  # type: ignore[var-annotated]
 _DISPATCHER = None
 
 
@@ -141,14 +141,14 @@ def make_custom_loop_for(
 
 
 def convert_custom_loop_into_onnx(
-    g: "GraphBuilder",  # noqa: F821
+    g: Any,  # "GreaphBuilder"
     sts: Dict[str, Any],
     outputs: List[str],
     *args: str,
     body: onnx.GraphProto,
     reduction_dim: Optional[Tuple[int, ...]] = None,
     name: str = "loop_for",
-) -> Tuple[str, ...]:
+) -> Union[str, Tuple[str, ...]]:
     graph = body.graph if isinstance(body, onnx.ModelProto) else body
     assert isinstance(
         graph, onnx.GraphProto

onnx_diagnostic/export/control_flow_research.py

@@ -29,7 +29,7 @@ def gen_schema(self, n_iter, body_fn, operands):
         from torch._higher_order_ops.schema import HopSchemaGenerator
         from torch._higher_order_ops.utils import materialize_as_graph
 
-        body_gm: torch.fx.GraphModule = materialize_as_graph(
+        body_gm: torch.fx.GraphModule = materialize_as_graph(  # type: ignore[annotation-unchecked]
             body_fn, (torch.tensor(0, dtype=torch.int64), *operands)
         )
         (