From e7d90a4b82be35ae7aed9bd801048205abe7de38 Mon Sep 17 00:00:00 2001 From: Scott Todd Date: Thu, 29 Feb 2024 13:01:13 -0800 Subject: [PATCH] [onnx] Fix type on create_module() in onnx_importer.py. (#2968) The type returned was changed in https://github.com/llvm/torch-mlir/pull/2795. This led to errors in the downstream IREE project: https://github.com/openxla/iree/pull/16622. --- python/torch_mlir/extras/onnx_importer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/torch_mlir/extras/onnx_importer.py b/python/torch_mlir/extras/onnx_importer.py index 91ee4c14c75a..289e5722efce 100644 --- a/python/torch_mlir/extras/onnx_importer.py +++ b/python/torch_mlir/extras/onnx_importer.py @@ -99,12 +99,12 @@ def __init__(self, model_proto: onnx.ModelProto, *, config: Config = Config()): assert model_proto.graph, "Model must contain a main Graph" self.main_graph = GraphInfo(self, model_proto.graph) - def create_module(self, context: Optional[Context] = None) -> Operation: + def create_module(self, context: Optional[Context] = None) -> Module: if not context: context = Context() - module_op = Module.create(Location.unknown(context)) + module = Module.create(Location.unknown(context)) # TODO: Populate module level metadata from the ModelProto - return module_op + return module class GraphInfo: