Hi @imnishantg
That’s an inconsistency with the batch shape.
When Fastai downloads a pretrained model, it expects 3 channels as input (RGB). Besides that, it adds to the Dataloader the after_batch Normalize transformation with the ImageNet stats.
That’s not true when we load a saved model or create a pretrained=False (as you did in the inference).
The dataloader with the after_batch normalization outputs 3 channels (example: [bs, 3, 224, 224]) and that fits the model. Your original dataloader outputs just 1 channel.
I wrote a detailed answer with example in the Kaggle notebook`s comments:
https://www.kaggle.com/cordmaur/fastai2-medical-simple-training/comments
Hope that helps.