Hi @JuanFMontesinos  and thanks for your answer.
Do you mean first load the state dict and then define the optimizer? I just tried but nothing.
Loaded model in training phase has zero gradient flow, basically it doesn’t train at all.