Unfortunately, I got different results:
I think something is wrong with batch normalization in UNet
After training model and before saving state_dict (model.train())
tensor([[[[ -8.4269, -12.0579, -11.8041,  …, -10.0886,  -8.6157,  -5.6419],
[ -5.6116, -10.0398,  -8.1797,  …,  -6.5734,  -5.6017,  -8.4666],
[ -7.3116, -12.3692, -10.6413,  …,  -7.6334,  -4.2797,  -6.6685],
…,
[ -3.6545,  -5.0984,  -5.1562,  …,  -3.8056,  -4.3882,  -9.5862],
[ -7.2274,  -4.3889,  -4.6838,  …,  -3.6184,  -3.2181, -13.1806],
[ -7.8890,  -4.6818,  -6.3442,  …,  -5.1584,  -4.0778,  -6.4020]]]],
device=‘cuda:0’, grad_fn=)
After saving state_dict and set model to evaluation mode (model.eval())
tensor([[[[-5.5660, -5.7427, -6.7510,  …, -6.1302, -6.1012, -7.4343],
[-5.5984, -5.3072, -4.8963,  …, -6.6118, -5.3268, -5.8644],
[-6.6022, -5.5419, -5.0204,  …, -6.9750, -4.1107, -5.6634],
…,
[-4.4239, -4.8906, -6.2160,  …, -5.3613, -4.3033, -3.7815],
[-6.3954, -5.0958, -6.3577,  …, -5.2069, -4.2653, -3.2605],
[-6.4229, -4.5915, -4.6084,  …, -3.9113, -3.7404, -3.7629]]]],
device=‘cuda:0’, grad_fn=)