More efficient way of indexing to avoid loop

vision
1

During my CNN training, I have:

  1. My image feature map with shape of [10, 256, 64, 64], where 10 is batch size and 256 is my channel, both 64’s are height and width of feature map size.

  2. Another tensor(points tensor) with shape of [10, 64, 2], which are basically stakes of 2d coordinates. So for each sample, I have 64 points with x and y coordinates.

My question is how to select the values on my feature map tensor based on the coordinate in my points tensor?
The output should ideally be in shape of [10, 64, 256], since for each point, I will have values from 256 channels.

How could I do that without running in a python loop?

My current approach is in below:

results = torch.zeros( 10, 64, 256)
feature_map = feature_map.transpose(1,2).transpose(2,3) # transpose feature map to shape of [10, 64, 64, 256]
for i in range(feature_map.shape[0]): # select i-th sample
    for j in range(points.shape[1]): # select j-th point for current sample
        value = feature_map[i][points[i][j][1]][points[i][j][0]] #get value at a certain coordinate
        results[i][j] = value

I am wondering if there will be a more efficient way of doing this kind of indexing. Any help will be appreciated.