Trainer.predict is not working on TPUs

[KijitoraButi]

🐛 Bug

Trainer.predict does not seem to run correctly when using TPU.
I get an error "TPUs do not currently support IterableDataset objects, the dataset must implement __len__. HINT: You can mock the length on your dataset to bypass this MisconfigurationException." even though I am using a dataset that implements __len__.

To Reproduce

If you create an instance of Trainer with tpu_cores=1 or 8, and run predict, an error occurs.

Notebook: https://colab.research.google.com/drive/1yHI-eKlfdq_zXUSPWmNvivEPJKl8RZbA?usp=sharing

Expected behavior

Return the output of the model.

Environment

• PyTorch Version (e.g., 1.0): 1.8
• Any other relevant information: Google Colab with TPU accelerator

[KijitoraButi]

Thanks for the fix!
However, I have a problem.
I tried the master branch and the error is no longer there, but none is returned as a result of predict.