diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 557b5d7d..b34d64f7 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -42,6 +42,7 @@ class ModelIO(object): class TrainerIO(object): def __get_model(self): + print(type(self.model)) is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel model = self.model.module if is_dp_module else self.model return model diff --git a/tests/README.md b/tests/README.md index fc60a6e2..0b747112 100644 --- a/tests/README.md +++ b/tests/README.md @@ -36,4 +36,11 @@ This file fits a tiny model on MNIST using these different set-ups. 3. Multiple (2) GPUs using DP + apex (for 16-bit precision). 3. Multiple (2) GPUs using DDP + apex (for 16-bit precision). +For each set up it also tests: +1. model saving. +2. model loading. +3. predicting with a loaded model. +4. simulated save from HPC signal. +5. simulated load from HPC signal. +