llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix call to fx_importer.module

fx_use_module_op
dan 2024-03-22 11:40:18 -05:00
parent d176528f4b
commit c75211f249
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
python/torch_mlir/fx.py
View File

@ -41,7 +41,7 @@ def export_and_import(
else: else:
fx_importer.import_frozen_program(prog, func_name=func_name) fx_importer.import_frozen_program(prog, func_name=func_name)
return fx_importer.module return fx_importer.module_op
def stateless_fx_import( def stateless_fx_import(
@ -55,4 +55,4 @@ def stateless_fx_import(
if fx_importer is None: if fx_importer is None:
fx_importer = FxImporter(context=context, hooks=hooks) fx_importer = FxImporter(context=context, hooks=hooks)
fx_importer.import_stateless_graph(gm.graph, func_name=model_name) fx_importer.import_stateless_graph(gm.graph, func_name=model_name)
return fx_importer.module return fx_importer.module_op

2
test/python/fx_importer/sparse_test.py
View File

@ -130,7 +130,7 @@ def export_and_import(f, *args, **kwargs):
fx_importer = FxImporter(context=context) fx_importer = FxImporter(context=context)
prog = sparse_export(f, args, kwargs) prog = sparse_export(f, args, kwargs)
fx_importer.import_frozen_program(prog) fx_importer.import_frozen_program(prog)
return fx_importer.module return fx_importer.module_op
def sparse_jit(f, *args, **kwargs): def sparse_jit(f, *args, **kwargs):