Change the "functional.interpolate" layer to "functional.max_unpool2d"

I have this model which has an encoder-decoder structure. I want to change the “functional.interpolate” layer (which does upsampling layer) to “functional.max_unpool2d”. But as you can see in the following code in the “colab link”, the max pooling is inside a Sequential, and I couldn’t get the indices of the tensor of max pooling layer. Are there any ideas on how to change the “functional.interpolate” layer to “functional.max_unpool2d” without changing a lot in the code. Really appreciate any help. Thank you so much in advance.

To use this layer in an nn.Sequential container, you could take a look at this post, which shows some approaches how multiple inputs can be used.

Thank you so much @ptrblck. That is very helpful.
May I ask you, what is the advantage that can I expect from using “functional.max_unpool2d” instead of “functional.interpolate”?. What I have understood from my reading, that will yield a cheaper computation model. Is this correct?

I don’t know the theoretical complexities of both algorithms and how they are implemented.
If you are concerned about the performance of these layers, you could profile them via e.g. the torch.utils.benchmark utilities.