update predict_step in ddp and ddp_spawn

Lightning-AI · Aug 3, 2021 · 6fe7e55 · 6fe7e55
1 parent 5af7ac4
commit 6fe7e55
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 8 deletions.
diff --git a/pytorch_lightning/plugins/training_type/ddp.py b/pytorch_lightning/plugins/training_type/ddp.py
@@ -399,10 +399,7 @@ def test_step(self, *args, **kwargs):
         return self.lightning_module.test_step(*args, **kwargs)
 
     def predict_step(self, *args, **kwargs):
-        if isinstance(self.model, DistributedDataParallel):
-            return self.model(*args, **kwargs)
-        else:
-            return self.lightning_module.predict_step(*args, **kwargs)
+        return self.lightning_module.predict_step(*args, **kwargs)
 
     def post_training_step(self):
         if not self.lightning_module.automatic_optimization:

diff --git a/pytorch_lightning/plugins/training_type/ddp_spawn.py b/pytorch_lightning/plugins/training_type/ddp_spawn.py
@@ -374,10 +374,7 @@ def test_step(self, *args, **kwargs):
         return self.lightning_module.test_step(*args, **kwargs)
 
     def predict_step(self, *args, **kwargs):
-        if isinstance(self.model, DistributedDataParallel):
-            return self.model(*args, **kwargs)
-        else:
-            return self.lightning_module.predict_step(*args, **kwargs)
+        return self.lightning_module.predict_step(*args, **kwargs)
 
     def post_training_step(self):
         if not self.lightning_module.automatic_optimization: