Simplify replicate_model_fn.GatheringOptimizer inspired by the past comments.
- I worked around the need to rely on Optimizer.__class__ for keeping track of the gradients. Now we are relying on the order of collecting them. I also added a basic verification that all towers have indeed called the same number of optimizers. - I now allow the user to increment global step however many times they wish. The changes above allowed me to support using the same optimizer class multiple times in a tower. I also renamed GatheringOptimizer to TowerOptimizer, which is a breaking change. #lifeincontrib PiperOrigin-RevId: 181381569
Loading
Please sign in to comment