@bibekx That’s a bit weird. I can’t tell how a variable change may have caused that.
What I would suggest is:

Check the model weights in the test_model function to ensure that an updated model is in fact passed.
Check that the dataloader in test_model is iterating through correctly.
Play around with the optimizer params (and different optimizers) and epochs. The loss flatlines at only the 4th epoch.

Personally, I don’t think 1 or 2 is an issue. It has to be the optimizer. Also, before any retraining, run cells 6, 10 and 11 just to ensure all the weights get reinitialized.