Solved by ptrblck in post #4 
              
              
                After you’ve stored the state_dicts you could iterate the keys of them and create a new state_dict using the mean (or any other reduction) for all parameters. 
This code snippet shows a small example: 
# Setup
state_dicts = []
model = models.resnet18()
optimizer = torch.optim.SGD(model.parameters(),…