frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-02 14:34:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update

Browse Source 
[ghstack-poisoned]
This commit is contained in:
Xuehai Pan
2024-10-16 16:34:57 +08:00
parent c19ca5328a a1a75d3c22
commit 134d27cf9d
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torch/utils/_cxx_pytree.py
View File

@ -866,6 +866,8 @@ def _broadcast_to_and_flatten(
def treespec_dumps(treespec: TreeSpec, protocol: Optional[int] = None) -> str:
"""Serialize a treespec to a JSON string."""
if 'treespec' not in type(treespec).__name__:
raise TypeError("treespec_dumps expects a treespec object.")
dummy_tree = tree_unflatten([0] * treespec.num_leaves, treespec)
orig_treespec = python_pytree.tree_structure(dummy_tree)
return python_pytree.treespec_dumps(orig_treespec, protocol=protocol)