Given two tensors of the same size, how can we use the indices obtained from max_pooling one tensor to subsample or pool the other tensor in PyTorch? When attempting this approach,
import torch
# Assuming y1 and e1 have the same dimensions
y1 = torch.randn(1, 32, 28, 28)
e1 = torch.randn(1, 32, 28, 28)
# Max-pooling
mxp = torch.nn.MaxPool2d(2, stride=2, return_indices=True)
y1_pooled, idx = mxp(y1)
# Subsample e1 using idx
e1 = e1[idx]
I get the error:
Traceback (most recent call last): e1 = e1[idx] IndexError: index 1 is out of bounds for dimension 0 with size 1