I found the problem. It was in upsampling of the output from the layers:
def get_layer8(self, input, output):
global output_layer8
upsampling_layer8 = nn.Upsample(scale_factor=2, mode=‘nearest’)
output_layer8 = upsampling_layer8(output.data)
output_layer8 = output_layer8.data
When I don’t upsample the tensor it works fine. I just have to manually get the correct pixels from the tensor which gets smaller after each layer.