including weights in .onnx

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>
borisfom · May 8, 2024 · d3c41f7 · d3c41f7
1 parent 53102bc
commit d3c41f7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/nemo/core/classes/common.py b/nemo/core/classes/common.py
@@ -1005,7 +1005,7 @@ def __init__(
         self.ignore_collections = ignore_collections
 
     def __call__(self, wrapped):
-        return self.unwrapped_call(wrapped) if is_typecheck_enabled() else self.unwrapped_call(wrapped)
+        return self.wrapped_call(wrapped) if is_typecheck_enabled() else self.unwrapped_call(wrapped)
 
     def unwrapped_call(self, wrapped):
         return wrapped

diff --git a/nemo/core/classes/exportable.py b/nemo/core/classes/exportable.py
@@ -222,10 +222,12 @@ def _export(
                         dynamic_axes.update(get_dynamic_axes(self.output_module.output_types_for_export, output_names))
                     if use_dynamo:
                         options = torch.onnx.ExportOptions(dynamic_shapes=dynamic_axes)
-                        ex_model = torch.export.export(jitted_model, tuple(input_list), kwargs=input_dict)
+                        ex_model = torch.export.export(
+                            jitted_model, tuple(input_list), kwargs=input_dict, strict=False
+                        )
                         ex_model = ex_model.run_decompositions()
                         ex = torch.onnx.dynamo_export(ex_model, *input_list, **input_dict, export_options=options)
-                        ex.save(output)
+                        ex.save(output, model_state=jitted_model.state_dict())
                         input_names = None
                     else:
                         torch.onnx.export(