This should work:
x = torch.tensor([[[0, 11, 12],
                   [3, 41, 51]],
                  [[10, 1, 21],
                   [31, 4, 51]],
                  [[40, 14, 2],
                   [34, 44, 5]]])

out = x[torch.arange(x.size(0)), :, [0, 1, 2]]
print(out)
> tensor([[0, 3],
          [1, 4],
          [2, 5]])