@JuanFMontesinos How should I specify the size for interpolate? For example, the output of up2 layer is  #torch.Size([1, 128, 120, 160]) so, it does have 128 channels, and batch is one. the question is what is should be size of input for interpolate to reach to the input original size which was [240, 320]?
#torch.Size([1, 128, 120, 160])
imgs = activations['up2'].squeeze().split(1, 0) 
           
imgs = torch.nn.functional.interpolate(imgs, size=None, mode='bilinear')