Right now, whenever save_checkpoint() is called, the EMA parameters are copied to the UNet. However, the original parameters are not restored, meaning that training continues on the EMA parameters, which is undesired and means that training would be impacted by the frequency of args.save_steps.
Added store and restore functions to EMAModel to be able to retain the original parameters; these functions are pulled from the original CompVis code.
Right now, whenever
save_checkpoint()is called, the EMA parameters are copied to the UNet. However, the original parameters are not restored, meaning that training continues on the EMA parameters, which is undesired and means that training would be impacted by the frequency ofargs.save_steps.Added
storeandrestorefunctions toEMAModelto be able to retain the original parameters; these functions are pulled from the original CompVis code.