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