| 28 | return x |
| 29 | |
| 30 | class CausalBlock1D(Block1D): |
| 31 | def __init__(self, dim: int, dim_out: int): |
| 32 | super(CausalBlock1D, self).__init__(dim, dim_out) |
| 33 | self.block = torch.nn.Sequential( |
| 34 | CausalConv1d(dim, dim_out, 3), |
| 35 | Transpose(1, 2), |
| 36 | nn.LayerNorm(dim_out), |
| 37 | Transpose(1, 2), |
| 38 | nn.Mish(), |
| 39 | ) |
| 40 | |
| 41 | def forward(self, x: torch.Tensor, mask: torch.Tensor): |
| 42 | output = self.block(x * mask) |
| 43 | return output * mask |
| 44 | |
| 45 | |
| 46 | class CausalResnetBlock1D(ResnetBlock1D): |