Skip to content

Commit 94f54e8

Browse files
angelayifacebook-github-bot
authored andcommitted
Add custom treespec fqn field
Summary: X-link: pytorch/pytorch#112428 Custom classes that are serialized with pytree are serialized by default with `f”{class.__module__}.{class.__name__}”`. This is a dependency from our serialized program directly into the outer Python environment. If a user moves the class to a different directory, the serialized program will be unable to be loaded. So, we will require users to pass in an FQN if they want to serialize their custom treespec type. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng avikchaudhuri gmagogsfm zhxchen17 tugsbayasgalan suo imported-using-ghimport Reviewed By: suo Differential Revision: D50886366 Pulled By: angelayi fbshipit-source-id: 276af112307e3d89137ff10af999ebeb5935a5ad
1 parent bd6a766 commit 94f54e8

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

userbenchmark/dynamo/dynamobench/common.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -566,7 +566,10 @@ def _register_dataclass_output_as_pytree(example_outputs) -> None:
566566
for output_type in output_dataclass_types:
567567
from torch._export.utils import register_dataclass_as_pytree_node
568568

569-
register_dataclass_as_pytree_node(output_type)
569+
register_dataclass_as_pytree_node(
570+
output_type,
571+
serialized_type_name=f"{output_type.__module__}.{output_type.__name__}",
572+
)
570573

571574

572575
class Stats:

0 commit comments

Comments
 (0)