Right now, the objax.optimizer.ExponentialMovingAverage is very generic as it only focuses on variables. Make an additional version objax.optimizer.ExponentialMovingAverageModule that handles 99% of use cases: take a model and return an EMA model of it.
Right now, the
objax.optimizer.ExponentialMovingAverageis very generic as it only focuses on variables. Make an additional versionobjax.optimizer.ExponentialMovingAverageModulethat handles 99% of use cases: take a model and return an EMA model of it.