From c75211f249532359946d7c9bc4ba52d8d81cad16 Mon Sep 17 00:00:00 2001 From: dan Date: Fri, 22 Mar 2024 11:40:18 -0500 Subject: [PATCH] fix call to fx_importer.module --- python/torch_mlir/fx.py | 4 ++-- test/python/fx_importer/sparse_test.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/torch_mlir/fx.py b/python/torch_mlir/fx.py index 298c66e26..8df43623e 100644 --- a/python/torch_mlir/fx.py +++ b/python/torch_mlir/fx.py @@ -41,7 +41,7 @@ def export_and_import( else: fx_importer.import_frozen_program(prog, func_name=func_name) - return fx_importer.module + return fx_importer.module_op def stateless_fx_import( @@ -55,4 +55,4 @@ def stateless_fx_import( if fx_importer is None: fx_importer = FxImporter(context=context, hooks=hooks) fx_importer.import_stateless_graph(gm.graph, func_name=model_name) - return fx_importer.module + return fx_importer.module_op diff --git a/test/python/fx_importer/sparse_test.py b/test/python/fx_importer/sparse_test.py index 52f10de32..df8e6980c 100644 --- a/test/python/fx_importer/sparse_test.py +++ b/test/python/fx_importer/sparse_test.py @@ -130,7 +130,7 @@ def export_and_import(f, *args, **kwargs): fx_importer = FxImporter(context=context) prog = sparse_export(f, args, kwargs) fx_importer.import_frozen_program(prog) - return fx_importer.module + return fx_importer.module_op def sparse_jit(f, *args, **kwargs):