Solved by Eta_C in post #2 
              
              
                import torch
import torch.nn as nn

A=torch.tensor([[5, 2, 6], [7, 3, 4]]).float()  
B=torch.tensor([[0, 2, 1, 1], [2, 2, 1, 0]]).long()
result = torch.gather(A, 1, B)
print(result)