denishem:

I didn’t quite understand your suggestion, should I save the images that I use for training and run inference on them for comparison? For the training part I am using DataLoader for both training and test datasets.


This would be one possibility to further isolate it.
Issues like these are often caused by either a different data loading and processing pipeline (majority of the issues) or by an invalid model loading (e.g. using strict=False as a “workaround” to load an invalid state_dict).
To further isolate the root cause I would suggest to start with the data loading and verify that the training and test script create the “same” input data. Note that the tensors wouldn’t be exactly the same, if the training pipeline uses random transformations, so for debugging purposes you could disable them.
The model can be checked by comparing the output for a static input (e.g. torch.ones) and make sure both scripts return an allclose output.