| 3 | |
| 4 | |
| 5 | class SeqToANNContainer(nn.Module): |
| 6 | def __init__(self, *args): |
| 7 | super().__init__() |
| 8 | if len(args) == 1: |
| 9 | self.module = args[0] |
| 10 | else: |
| 11 | self.module = nn.Sequential(*args) |
| 12 | |
| 13 | def forward(self, x_seq: torch.Tensor): |
| 14 | y_shape = [x_seq.shape[0], x_seq.shape[1]] |
| 15 | y_seq = self.module(x_seq.flatten(0, 1).contiguous()) |
| 16 | y_shape.extend(y_seq.shape[1:]) |
| 17 | return y_seq.view(y_shape) |
| 18 | |
| 19 | |
| 20 | class SpikeModule(nn.Module): |