Solved by Eta_C in post #2 
              
              
                Try this 
import torch
A = torch.tensor([[0, 1], [2, 3], [5, 6]]).float()
B = torch.tensor([[1], [1], [0]]).long()
print(torch.gather(A, 1, B))