Add serialized type name to pytrees

huggingface · Dec 14, 2023 · 0ee9678 · 0ee9678
1 parent 3060899
commit 0ee9678
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py
@@ -310,6 +310,7 @@ def __init_subclass__(cls) -> None:
                 cls,
                 _model_output_flatten,
                 _model_output_unflatten,
+                serialized_type_name=f"{cls.__module__}.{cls.__name__}",
             )
 
     def __init__(self, *args, **kwargs):
@@ -442,6 +443,7 @@ def _model_output_unflatten(values: Iterable[Any], context: "_torch_pytree.Conte
         ModelOutput,
         _model_output_flatten,
         _model_output_unflatten,
+        serialized_type_name=f"{ModelOutput.__module__}.{ModelOutput.__name__}",
     )