This would be better:
def load_params(flattened):
    for p, avg_p in zip(model.parameters(), flattened):
        p.data.copy_(avg_p)

Also, note that they’re no longer flattened, so you might want to change the name.