Solved by Shisho_Sama in post #3 
              
              
                Thank you very much, I did save it the way you mentioned, but for resuming I used : 
scheduler.load_state_dict(checkpoint['scheduler']) 
and thats all, no need to change anything else. 
So basically this is how the resume section is defined now : 
 # optionally resume from a checkpoint
  if resume:
…