Solved by iffiX in post #2 
              
              
                The first evident error in your code is: 
msd = model.state_dict()
for name, param in model.named_parameters():
    msd[name].grad = conn.tensorget(f'{name}_grad')
model.load_state_dict(msd)

Since load_state_dict will not load the gradient, but only load registered nn.Parameter() and buffers of mod…