TypeError caused due to arguments not found at the end of TPU training
🐛 Bug

When using TPU for training if you pass arguments when initializing the model, the training runs fine and completes all the epochs but at the end lightning for some reason tries to load the checkpoint back, here when loading the arguments are not found for some reason causing the error.
Please reproduce using the BoringModel
Reproduced in the BoringModel, with lr passed during model initialization.

https://colab.research.google.com/drive/1CJ1WFQBvfe0-1oFV6BYl-8cAmvzjQQl_?usp=sharing
To Reproduce
Just change the runtime to TPU and run the notebook.

Expected behavior
Load the arguments passed before properly.

Environment
Tested on both Colab and Kaggle.