Fix flax summarization example: save checkpoint after each epoch and push checkpoint to the hub
# What does this PR do?

Current `run_summarization_flax.py` saves checkpoint & pushs to the hub UNDER the block `if training_args.do_predict:`

https://github.com/huggingface/transformers/blob/7051b892671982717388dd6dfca368a90da60252/examples/flax/summarization/run_summarization_flax.py#L811

This block is also outside the epoch looping block, and the saving/pushing won't happen after each epoch as expected.

This PR moves saving/pushing block to the right place.

## Who can review?

@patil-suraj  @patrickvonplaten 