Regression Not Working

[tahsin314]

I'm working on a regression task. My model is supposed to take an image of retina as input and predict the severity level of an abnormality. The severity level can be any integer between 0 and 4. Here is my lightning module. The code works perfectly if I use ordinal regression but does not converge when I use regression and MSELoss. (Train Loss and Validation Loss). I've tried varying learning rate from 10^-5 to 10^1 but there isn't any improvement. I've also tried to overfit my model but it didn't help either. It would be a great help if someone can look into it. Thanks.