fix type

Lightning-AI · Dec 4, 2021 · 3d6fa8b · 3d6fa8b
1 parent 0e0b6ae
commit 3d6fa8b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/plugins/training_type/ddp.py b/pytorch_lightning/plugins/training_type/ddp.py
@@ -415,7 +415,7 @@ def test_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
             return self.lightning_module.test_step(*args, **kwargs)
 
     def predict_step(self, *args, **kwargs) -> STEP_OUTPUT:
-        with self.precision_plugin.test_step_context():
+        with self.precision_plugin.predict_step_context():
             return self.lightning_module.predict_step(*args, **kwargs)
 
     def post_training_step(self):