Fixes #2455

Lightning-AI · Jul 2, 2020 · 80988a3 · 80988a3
1 parent 0264783
commit 80988a3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/callbacks/early_stopping.py b/pytorch_lightning/callbacks/early_stopping.py
@@ -149,8 +149,8 @@ def _run_early_stopping_check(self, trainer, pl_module):
             should_stop = self.wait_count >= self.patience
 
             # check flag across all GPUs
-            should_stop = torch.tensor(should_stop)
             if trainer.use_ddp or trainer.use_ddp2:
+                should_stop = torch.tensor(should_stop, device=pl_module.device)
                 dist.all_reduce(should_stop, op=dist.ReduceOp.MAX)
 
             # do actual stop