Solved by Mona_Jalal in post #5 
              
              
                I set strict to False here and it is working. I didn’t change anything else. 
tnet.load_state_dict(checkpoint['state_dict'], strict=False)