[export] Add serialized_type_name to torch.Size

2025-11-01 13:34:57 +08:00 · 2024-11-18 13:56:50 -08:00
1 changed files with 1 additions and 0 deletions
--- a/torch/fx/experimental/proxy_tensor.py
+++ b/torch/fx/experimental/proxy_tensor.py
@ -131,6 +131,7 @@ pytree.register_pytree_node(
        [(pytree.SequenceKey(i), x) for i, x in enumerate(xs)],
        None,
    ),
+    serialized_type_name="torch.Size",
 )