Update weights of the average model with new model observation

def update_average_model(self, model):
        """ Update weights of the average model with new model observation """
        for model_param, average_param in zip(model.parameters(), self.average_model.parameters()):
            # EWMA average model update
            average_param.data.mul_(self.average_model_alpha).add_(model_param.data * (1 - self.average_model_alpha))