From a3ad0e0ac1dbc27642aa1f10f17302905f4fe03e Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 19:23:11 -0400
Subject: [PATCH] ignoring dist parallel forward

---
 pytorch_lightning/root_module/root_module.py | 1 -
 tests/test_models.py                         | 8 ++++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
index 7f99ef98..b2a08830 100644
--- a/pytorch_lightning/root_module/root_module.py
+++ b/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
diff --git a/tests/test_models.py b/tests/test_models.py
index e4ca171b..fd5fcf51 100644
--- a/tests/test_models.py
+++ b/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()