I need to implement a pooling layer, which will pool from a given tensor, based on the indices generated by the max pooling on another tensor. For example,
import torch
import torch.nn as nn

# Define a tensor
X = torch.rand(2, 3, 4, 4)
X[0, 0, :, :]                                                                                                                                                                         
tensor([[0.9889, 0.2736, 0.6786, 0.8688],
        [0.5038, 0.4131, 0.8648, 0.4014],
        [0.1842, 0.8754, 0.5671, 0.6143],
        [0.1987, 0.8137, 0.6612, 0.5943]])

# Define another tensor
Y = torch.rand(2, 3, 4, 4)
Y[0, 0, :, :]                                                                                                                                                                         
tensor([[0.7160, 0.2788, 0.1997, 0.4487],
        [0.0241, 0.2625, 0.3674, 0.5985],
        [0.1600, 0.4911, 0.1724, 0.0827],
        [0.6873, 0.8463, 0.5221, 0.8582]])

# Define a maximum pooling layer        
max_pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        
# Apply maximum pooling on tensor `X`
X_p = max_pool(X)
X_p[0, 0, :, :]                                                                                                                                                                      
tensor([[0.9889, 0.8688],
        [0.8754, 0.6612]])

Now, what I would like to do is to pool from tensor Y using the indices of the maximum values of tensor X. The pooling result on tensor Y should be the following:
Y_p[0, 0, :, :]                                                                                                                                                                      
tensor([[0.7160, 0.4487],
        [0.4911, 0.5221]])

Thank you!