Pytorch index for 3D-tensor

Hello, I have a 3D tensor A of size (x,y,z) and an 1D vector B of size (y), in which each (index,value) in B correspond to a pair of index of y-dimension and z-dimension. The result slice tensor will have dimension (x,y) or (x,y,1). Is there a fast way in Pytorch so that I can get the index that can be used to slice the tensor? Thanks a lot!

Could you post a dummy example using some random input tensors, as I’m unsure how to interpret the given shapes?