Fitting hangs at "cleaning up ddp environment..." when tpu_cores=8

[KijitoraButi]

🐛 Bug

When setting tpu_cores of Trainer to 8, fitting hangs at "cleaning up ddp environment...".

Please reproduce using the BoringModel

https://colab.research.google.com/drive/1tJswNaT0I-GrGsi6ngwwRDUmeFY1pFr3?usp=sharing

To Reproduce

Run above URL notebook.

Expected behavior

Trainer.fit ends normally .

Environment

• PyTorch Version (e.g., 1.0): 1.7.0a0+7e71a98
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source): N/A
• Python version: 3.7.10
• CUDA/cuDNN version: N/A
• GPU models and configuration: N/A
• Any other relevant information: I tried the notebook on Google Colab with TPU.

[kaushikb11]

Hi @KijitoraButi, thank you for opening the issue! Yes, tried to reproduce the issue & it does hang. We will look into fixing it soon.

[KijitoraButi]

I was able to train the model properly.
Thanks for the fix!

[kaushikb11]

@KijitoraButi Awesome! Did you try the master?

[KijitoraButi]

@kaushikb11
Yes!
I updated Lightning to Master and then ran the code and it worked fine.