I would like to use the running average of parameters instead of using the parameters from the training directly at the test session.
To do this, I initialized the running average parameters from the network as
avg_param = torch.cat([param.view(-1) for param in model.parameters()],0)

Then, I performed the running average at each training iteration as
avg_param = 0.9*avg_param + 0.1*torch.cat([param.data.view(-1) for param in model.parameters()],0)

Finally, at the test session, I loaded the parameters as
i = 0
for param in model.parameters():
    param = avg_param[i:i+param.nelement()].resize(*param.size())
    i = i+param.nelement()

Is this process correct ?