torch tensor.permute()

torch.Tensor.transpose() swaps two dimensions, while permute() is a generalized operation that allows you to specify the order of all axes