Does advanced indexing (a[:, b]) work for you (note that you need a byte/uint8 tensor to index)?
a = torch.ones((3, 3, 3))
a[:, 1, 1] = 2
a[:, 2, 2] = 5

b = torch.zeros((3, 3), dtype=torch.uint8)
b[1, 1] = 1
b[2, 2] = 1

print(a[:, b])

Best regards
Thomas