Best model checkpoints only save rank zero part of the pipe paralleled model.
🐛 Best model checkpoints only save rank zero part of the pipe paralleled model.
When saving the best checkpoint of the piped parallel model in the relevant module, only a part of the model is saved. The method of storing rank zero is suitable for the method of DP, but there is an issue in the MP.
I confirmed that the code below works well.
haven-jeon@9ebe8c0#diff-05daced8d811d9d4566136e3db8c0c35dac3312e6ca6ee6ba0ef7678fa16d051
Please reproduce using the BoringModel
To Reproduce
Use following BoringModel and post here

Expected behavior

Environment
Note: Bugs with code are solved faster ! Colab Notebook should be made public !


IDE: Please, use our python bug_report_model.py template.


Colab Notebook: Please copy and paste the output from our environment collection script (or fill out the checklist below manually).


You can get the script and run it with:
wget https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py


PyTorch Version (e.g., 1.0): 1.6.0
OS (e.g., Linux): Centos 7
How you installed PyTorch (conda, pip, source): pip
Build command you used (if compiling from source):
Python version: 3.7
CUDA/cuDNN version: 10.2
GPU models and configuration: P40
Any other relevant information:

Additional context