| 10 | |
| 11 | class SharedMLP(nn.Module): |
| 12 | def __init__(self, in_channels, out_channels, dim=1): |
| 13 | super().__init__() |
| 14 | if dim == 1: |
| 15 | conv = nn.Conv1d |
| 16 | bn = nn.GroupNorm |
| 17 | elif dim == 2: |
| 18 | conv = nn.Conv2d |
| 19 | bn = nn.GroupNorm |
| 20 | else: |
| 21 | raise ValueError |
| 22 | if not isinstance(out_channels, (list, tuple)): |
| 23 | out_channels = [out_channels] |
| 24 | layers = [] |
| 25 | for oc in out_channels: |
| 26 | layers.extend([ |
| 27 | conv(in_channels, oc, 1), |
| 28 | bn(8, oc), |
| 29 | Swish(), |
| 30 | ]) |
| 31 | in_channels = oc |
| 32 | self.layers = nn.Sequential(*layers) |
| 33 | |
| 34 | def forward(self, inputs): |
| 35 | if isinstance(inputs, (list, tuple)): |