move pass

Signed-off-by: Ali Boubezari <[email protected]>

Author: aboubezari

modelopt/onnx/autocast/graphsanitizer.py

@@ -63,9 +63,9 @@ def sanitize(self) -> None:
         self.ensure_graph_name_exists()
         onnx_utils.name_onnx_nodes(self.model.graph)
         self.replace_custom_domain_nodes()
+        self.sanitize_io_casts()
         self.cleanup_model()
         self.set_ir_version(self.max_ir_version)
-        self.sanitize_io_casts()
 
     def find_custom_nodes(self) -> None:
         """Find custom nodes in the model.
@@ -356,6 +356,10 @@ def sanitize_io_casts(self) -> None:
         for node in nodes_to_add:
             self.model.graph.node.append(node)
 
+        # Make sure the graph is topologically sorted
+        gs_graph = gs.import_onnx(self.model).cleanup().toposort()
+        self.model = gs.export_onnx(gs_graph)
+
     def _create_layernorm_node(self, pattern: dict) -> onnx.NodeProto:
         """Create a LayerNormalization node with optional bias."""
         ln_name = f"LayerNorm_{pattern['mean_node'].name}"