(self, dim0: int, dim1: int)
| 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) |