Solved by sxbio in post #3 
              
              
                Just to confirm that the weight decay may have been the culprit… Added the following snippet and things are training smoothly… 
def group_weight(module): 
group_decay = [] 
group_no_decay = [] 
for n,m in module.named_parameters(): 
if ‘input_embedding’ in n: 
group_no_decay.append(m) 
else: 
group_…