I got some answer related to this question. I don’t remember the source but thanks whoever has given the answer.
b,c,h,w = x.size() # let the size be(2,128,32,32)
slice_idx = [[i] for i in range(b)]
#idx is a tensor of size (2,32) containg the the indices to keep from 2nd dimension (c)
#the output should have size (2,32,32,32)
out= x[slice_idx, idx]

I want to ask whether what I’m doing Is correct?