Skip to content

Commit

Permalink
Remove redundant serialization code (#126249)
Browse files Browse the repository at this point in the history
After #123308, we no longer need separate serialization path to handle different types that exist in the `nn_module` metadata. This PR cleans up the redundant code.
Pull Request resolved: #126249
Approved by: https://github.com/angelayi
  • Loading branch information
jiashenC authored and ZelboK committed May 19, 2024
1 parent 2a34465 commit 45a699a
Showing 1 changed file with 2 additions and 14 deletions.
16 changes: 2 additions & 14 deletions torch/_export/serde/serialize.py
Original file line number Diff line number Diff line change
Expand Up @@ -539,21 +539,9 @@ def export_nn_module_stack(val):
path, ty = val

assert isinstance(path, str)
assert isinstance(ty, str)

# node.meta["nn_module_stack"] could have two forms:
# 1. (path: str, module_type: 'type'), e.g.
# ('', <class 'sigmoid.inference.MySimpleModel'>)
# 2. (path: str, module_type: str), e.g.
# ('', 'sigmoid.inference.MySimpleModel')
# ExportedProgram directly produced by torch.export() has form 1
# ExportedProgram deserialized from disk has form 2
# TODO: This is not ideal, we should fix this.
if isinstance(ty, str):
normalized_ty = ty
else:
normalized_ty = ty.__module__ + "." + ty.__qualname__

return path + "," + normalized_ty
return path + "," + ty

# Serialize to "key,orig_path,type_str"
nn_module_list = [
Expand Down

0 comments on commit 45a699a

Please sign in to comment.