ignoring dist parallel forward

pytorch_lightning/root_module/root_module.py

@@ -92,7 +92,6 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         model_summary = ModelSummary(self)
         print(model_summary)
 
-
     def freeze(self):
         for param in self.parameters():
             param.requires_grad = False

tests/test_models.py

@@ -43,6 +43,10 @@ def test_early_stopping_cpu_model():
     model, hparams = get_model()
     run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
 
+    # test freeze on cpu
+    model.freeze()
+    model.unfreeze()
+
 
 def test_cpu_model_with_amp():
     """
@@ -133,6 +137,10 @@ def test_amp_gpu_ddp_slurm_managed():
     trainer.hpc_save(save_dir, exp)
     trainer.hpc_load(save_dir, on_gpu=True)
 
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
     clear_save_dir()