Update outputs.py

huggingface · Jan 8, 2025 · f8b292d · f8b292d
1 parent 5655b22
commit f8b292d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/diffusers/utils/outputs.py b/src/diffusers/utils/outputs.py
@@ -65,12 +65,14 @@ def __init_subclass__(cls) -> None:
                     cls,
                     torch.utils._pytree._dict_flatten,
                     lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)),
+                    serialized_type_name=f"{cls.__module__}.{cls.__name__}"
                 )
             else:
                 torch.utils._pytree.register_pytree_node(
                     cls,
                     torch.utils._pytree._dict_flatten,
                     lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)),
+                    serialized_type_name=f"{cls.__module__}.{cls.__name__}"
                 )
 
     def __post_init__(self) -> None: