def METHOD_NAME(self) -> None:
    net = Net().to(self.device)
    input = torch.Tensor([range(10)]).to(self.device)
    self.assertEqual(
        net.permute_pooled_embeddings(input).view(10).tolist(),
        [6, 7, 8, 9, 0, 1, 5, 2, 3, 4],
    )