Commit a2039aa9 authored by Igor Saprykin's avatar Igor Saprykin Committed by TensorFlower Gardener
Browse files

Support multiple loss and multiple optimizers in replicate_model_fn. 

Instead of supplying `optimizer_fn`, the user is now expected to wrap their optimizer in GatheringOptimizer.  The latter will gather gradients, reduce and apply them.

There can be multiple instances of GatheringOptimizer inside the model.

PiperOrigin-RevId: 179899422
parent f1a1bd87
Show whitespace changes
Inline Side-by-side
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment