Fix Checkpoint issue when using Horovod distributed backend (PyTorchL…

…ightning#6947) Co-Authored-By: Adrian Wälchli <[email protected]>
Lightning-AI · Apr 11, 2021 · f10b2a4 · f10b2a4
1 parent fe0d088
commit f10b2a4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/plugins/training_type/horovod.py b/pytorch_lightning/plugins/training_type/horovod.py
@@ -150,7 +150,7 @@ def reduce(self, tensor, group: Optional[Any] = None, reduce_op: Optional[Union[
 
         if reduce_op in (None, "avg", "mean"):
             reduce_op = hvd.Average
-        elif reduce_op == "sum":
+        elif reduce_op in ("sum", ReduceOp.SUM):
             reduce_op = hvd.Sum
         else:
             raise ValueError(f"unrecognized `reduce_op`: {reduce_op}")