Wandb logger doesn't upload saved model checkpoint for final epoch
🐛 Bug
When training a model on the TPU and using the wandb logger, the checkpoint for the last epoch trained doesn't get uploaded to wandb.
To Reproduce
Colab notebook: https://colab.research.google.com/drive/1oPaRWGZcz6YEol012xFADN42LV-jowtT