I found the error. It was stemming from the fact that I was saving the optimizer after each epoch without calling optimizer.consolidate_state_dict() first. After I included that before saving the warning went away.