Found the root cause for my problem. Nothing to do with model saving or loading. Just a silly mistake.

  
    
    
    Unable to pause and resume training (loading a model) without getting a small jump in training loss 
  
  
    After spending few hours trying to figure out the problem, found the root cause. And it has nothing to do with model / optimizer loading and saving at all. 
What I needed to change is 
lab_to_val = {j:i for i,j in enumerate(training_classes)}

To 
lab_to_val = {j:i for i,j in enumerate(sorted(training_classes))}

Since the label to value conversion was not happening on the sorted list, each time my runtime was reset, a single label would take different values. Ex - if person class takes value 1 …