Yes, this should work:
batch_size, N = 4, 5
A = torch.randn(batch_size, N, N)
x, y = torch.randint(0, N, (batch_size, N)), torch.randint(0, N, (batch_size, N))

B = torch.zeros(batch_size, N, N)

# loop approach
for b in range(batch_size):
    B[b]=A[b,x[b].unsqueeze(1),y[b]]


# avoiding loops
C = A[torch.arange(A.size(0)).unsqueeze(1).unsqueeze(2), x.unsqueeze(2), y.unsqueeze(1)]

# check for equal values
print((B == C).all())
> tensor(True)