Done!
Problem was that I made changes expecting loss from aux in eval also. My bad. Thanks @jdhao 
For others who get same problem this link may help.
For devs it would be good if necessary if conditions are put in main.py for different models.
3 Likes