From 72db64d294d4e1bc69806f78a4eb6489c6c87c5f Mon Sep 17 00:00:00 2001 From: ananthsub Date: Fri, 4 Feb 2022 19:33:25 -0800 Subject: [PATCH] Use the strategy's `root_device` instead of the LightningModule's device property (#11734) --- pytorch_lightning/strategies/parallel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/strategies/parallel.py b/pytorch_lightning/strategies/parallel.py index 5840e7816f..11207065b7 100644 --- a/pytorch_lightning/strategies/parallel.py +++ b/pytorch_lightning/strategies/parallel.py @@ -93,7 +93,7 @@ class ParallelStrategy(Strategy, ABC): return all_gather_ddp_if_available(tensor, group=group, sync_grads=sync_grads) def reduce_boolean_decision(self, decision: bool) -> bool: - decision = torch.tensor(int(decision), device=self.lightning_module.device) + decision = torch.tensor(int(decision), device=self.root_device) decision = self.reduce(decision, reduce_op=ReduceOp.SUM) decision = bool(decision == self.world_size) return decision