@ptrblck thank you for your reply.
I am using the “O1” opt_level.
Regarding the linear layer, it depends but the difference can be huge. For instance the first linear layer, I can have value around 10e1 or 10e2 insteand of 10-e1 and for the last layer, it is worse : the output of the last layer can be around 10e3 or 10e4 instead of usually being around 10-1.