Fixing "Found an invalid max index" error in PyTorch's max_unpool2d function

Question

I'm having a hard time using F.max_unpool2d in PyTorch. I define a tensor x, perform max pooling on it using nn.MaxPool2d and store the output and indices in y and indices. Then, I try to unpool y using F.max_unpool2d.

Here is the minimal code:

import torch
from torch import nn
from torch.nn import functional as F

x = torch.randn(1, 64, 11, 11)

Max_pool = nn.MaxPool2d(2, stride=2, return_indices=True)

y, indices = Max_pool(x)

print(y.shape)
print(indices.shape)

# -- unpooling
y_ = F.max_unpool2d(y, indices, kernel_size=2, stride=2, padding=0)

print(y_.shape)

However, I get the following error:

max_unpool2d
    return torch._C._nn.max_unpool2d(input, indices, output_size)
RuntimeError: Found an invalid max index: 108 (output volumes are of size 10x10)