Register subclasses as pytree nodes. This is necessary to synchronize gradients when using `torch.nn.parallel.DistributedDataParallel` with `static_graph=True` with modules that output `ModelOutput` subclasses.
(cls)
| 48 | """ |
| 49 | |
| 50 | def __init_subclass__(cls) -> None: |
| 51 | """Register subclasses as pytree nodes. |
| 52 | |
| 53 | This is necessary to synchronize gradients when using `torch.nn.parallel.DistributedDataParallel` with |
| 54 | `static_graph=True` with modules that output `ModelOutput` subclasses. |
| 55 | """ |
| 56 | if is_torch_available(): |
| 57 | import torch.utils._pytree |
| 58 | |
| 59 | if is_torch_version("<", "2.2"): |
| 60 | torch.utils._pytree._register_pytree_node( |
| 61 | cls, |
| 62 | torch.utils._pytree._dict_flatten, |
| 63 | lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)), |
| 64 | ) |
| 65 | else: |
| 66 | torch.utils._pytree.register_pytree_node( |
| 67 | cls, |
| 68 | torch.utils._pytree._dict_flatten, |
| 69 | lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)), |
| 70 | serialized_type_name=f"{cls.__module__}.{cls.__name__}", |
| 71 | ) |
| 72 | |
| 73 | def __post_init__(self) -> None: |
| 74 | class_fields = fields(self) |
nothing calls this directly
no test coverage detected