Allow passing in a tensor to tfa.optimizers.MovingAverage num_updates

[Hyperparticle]

Currently, tfa.optimizers.MovingAverage has an assert that requires num_updates to be of type int, shown here. This prevents me from passing in an integer tensor that changes with the global step, which is officially supported by tf.train.ExponentialMovingAverage.

Can this assert be updated to handle this use case?

[SSaishruthi]

@seanpmorgan I was looking at the code and looks like we just need to remove the check for int. Is that right? I will create a PR if that is right

[Hyperparticle]

Yes, either remove the int check or additionally check if the input is a tensor.

[SSaishruthi]

Thanks. Creating PR now

[guillaumekln]

Is the num_updates argument necessary? Why not using optimizer.iterations directly?

[Hyperparticle]

It's possible you might want to change the moving average over time with custom behavior and num_updates would do that. However, the current code changes have broken this again. It looks like you won't be able to change the moving average ratio every step, only at the beginning.