>>> a = torch.randn(1, 2, 3, 4)
>>> a_flip = flip(a, [1, 3])
>>> a_flip.shape
torch.Size([2, 4, 1, 3]) # need to permute([2, 0, 3, 1])

# =====  original flip() impl only works for this case =======
>>> a = torch.randn(1, 2, 3, 4)
>>> a_flip = flip(a, [0, 2])
>>> a_flip.shape
torch.Size([1, 3, 2, 4]) # need to permute([0, 2, 1, 3])

>>> a = torch.randn(1, 2, 3, 4)
>>> a_flip = flip(a, [0, 3])
>>> a_flip.shape
torch.Size([1, 4, 2, 3]) # need to permute([0, 2, 3, 1])

>>> a = torch.randn(1000, 1000)
>>> %timeit -r 10 a.flip(0, 1)
15.2 ms ± 2.57 ms per loop (mean ± std. dev. of 10 runs, 100 loops each)