| 18 | from matcha.models.components.transformer import BasicTransformerBlock |
| 19 | |
| 20 | class Transpose(torch.nn.Module): |
| 21 | def __init__(self, dim0: int, dim1: int): |
| 22 | super().__init__() |
| 23 | self.dim0 = dim0 |
| 24 | self.dim1 = dim1 |
| 25 | |
| 26 | def forward(self, x: torch.Tensor): |
| 27 | x = torch.transpose(x, self.dim0, self.dim1) |
| 28 | return x |
| 29 | |
| 30 | class CausalBlock1D(Block1D): |
| 31 | def __init__(self, dim: int, dim_out: int): |