To simply average the parameters of multiple models you could use this approach. However, based on your description you would like to calculate the gradients for the parameters in the original models, so using the averaged state_dict wouldn’t work.
Instead I think the right approach would be to create the new averaged tensors from all models and use the functional API for these new tensors.
Something like this would work:
model1 = nn.Linear(10, 10)
model2 = nn.Linear(10, 10)
new_weight = (model1.weight + model2.weight) / 2.
new_bias = (model1.bias + model2.bias) / 2.
x = torch.randn(1, 10)
out = F.linear(x, new_weight, new_bias)
out.mean().backward()
for name, param in model1.named_parameters():
print(name, param.grad)
for name, param in model2.named_parameters():
print(name, param.grad)
Although it seems to be quite cumbersome, so maybe someone else would have a better and cleaner approach.