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

View File

@ -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

View File

@ -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):