nn.Parameter does not update after training the first fold

[albanD]

Is the optimizer here a pytorch provided optimizer?
In general this part of the code is good.