def TransposedModule(module):
    """Wrap a SequenceModule class to accept transposed parameter, handle state, absorb kwargs."""
    # https://stackoverflow.com/a/65470430/1980685
    @functools.wraps(module, updated=())
    class TransposedModule(module):
        def __init__(self, *args, transposed=False, **kwargs):
            super().__init__(*args, **kwargs)
            self.transposed = transposed

        def forward(self, x, state=None, **kwargs):
            if self.transposed: x = x.transpose(-1, -2)
            x, next_state = super().forward(x, state) # Don't use kwarg because nn.LSTM
            next_state = None if state is None else next_state
            if self.transposed: x = x.transpose(-1,-2)
            return x, next_state
    # https://stackoverflow.com/questions/5352781/how-to-set-class-names-dynamically
    # TransposedModule.__name__ = module.__name__ # functools wraps is better solution
    return TransposedModule
Wrap a SequenceModule class to accept transposed parameter, handle state, absorb kwargs.
analyze/convnexts4nd/src/models/sequence/base.py