Move to optimizer based EMA implementation (#5169)

* Move to optimizer

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Fix replacing weights

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Allow swapping of weights be optional

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Save 2 models

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Use different hook

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Expose cpu device

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Add clause to see if this fixes issue with O2 optimizer

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Try to get O2 working

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* WIP

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Fixes

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Fixes to tests

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Add guard

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Remove import

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Add guard

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Add comment

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Remove overwrite

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Add BatchNorm, currently tests fail

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Fix tests/functionality for batch norm

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

* Get rid of NLP changes

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>

Signed-off-by: SeanNaren <snarenthiran@nvidia.com>