It’s certainly possible to create a Module instead, I’m dynamically creating the module contents so figured a sequential would be a nice fit. As an artificial example (it makes no sense to have activation=None here…):
class LstmOutput(nn.Sequential):
def __init__(self, lstm, activation=None):
sequence = [lstm]
if activation is not None:
sequence.append(activation)
super().__init__(*sequence)
self.lstm = lstm
def zero_state(self):
self.lstm.zero_state()
Of course, there are many ways to get around this: use a collections.OrderedDict so that the lstm attribute is added with the correct attribute name, use the apply function to zero_state the lstm, etc. Still, unless you inspect network architecture, it’s easy to miss that the order of operations is lstm -> activation -> lstm.