| 523 | |
| 524 | |
| 525 | class BatchNorm1dTBC(nn.Module): |
| 526 | def __init__(self, c): |
| 527 | super(BatchNorm1dTBC, self).__init__() |
| 528 | self.bn = nn.BatchNorm1d(c) |
| 529 | |
| 530 | def forward(self, x): |
| 531 | """ |
| 532 | |
| 533 | :param x: [T, B, C] |
| 534 | :return: [T, B, C] |
| 535 | """ |
| 536 | x = x.permute(1, 2, 0) # [B, C, T] |
| 537 | x = self.bn(x) # [B, C, T] |
| 538 | x = x.permute(2, 0, 1) # [T, B, C] |
| 539 | return x |
| 540 | |
| 541 | |
| 542 | class EncSALayer(nn.Module): |