Use xm.save to save model on TPU

[lezwon]

What does this PR do?

Fixes #2700

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together? Otherwise, we ask you to create a separate PR for every change.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Did you verify new and existing tests pass locally with your changes?
• If you made a notable change (that affects users), did you update the CHANGELOG?

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[pep8speaks]

Hello @lezwon! Thanks for updating this PR.

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2020-09-19 09:57:09 UTC

[edenlightning]

Hey @lezwon, any ETA on this fix?

[lezwon]

@edenafek it has a lot of breaking changes. Also there was an issue due to which it was hanging at the 4th epoch.
Had it on pause due to the refactors happening. Will try to get this done on the weekend. :)

[Borda]

@lezwon how is this going, can we get it done today?
pls mind rebase on master... maybe it would be easier for you to squash it fist and then rebase...

[edenlightning]

@lezwon can we revive this fix?

[lezwon]

@edenlightning I'm using a new branch to fix this issue. Will close this one as it is stale.

[Borda]

@lezwon mind link here the new PR?