how to update ema variables in different models, such as a big model and a small model ;

[qiuhui1991]

how to update ema variables in different models, such as a big model and a small model ;

[qiuhui1991]

def update_ema_variables(model, ema_model, global_step):
alpha = min(1 - 1 / (global_step + 1), 0.999)
for ema_param, param in zip(ema_model.parameters(), model.parameters()):
ema_param.data.mul_(alpha).add_(1 - alpha, param.data)

in different model, the function update_ema_variables produced error

[Vegeta2020]

@qiuhui1991 Can you give the traceback for your produced error? I guess it may be related with the settings of model/ema_model.