Fixes bug where "current_model" is unbound when task is restarted on final round.
Fixes bug where "current_model" is unbound when task is restarted on final round.

PiperOrigin-RevId: 349305478
Change-Id: I8f8262d3b7fc49f4a845698c74910ac5b58ac911

google-research/federated,google-research/federated