Fix Checkpoint issue when using Horovod distributed backend (PyTorchL…

…ightning#6947) (#6958) Co-Authored-By: Adrian Wälchli <[email protected]> Co-authored-by: Adrian Wälchli <[email protected]> (cherry picked from commit b37b58a)
Lightning-AI · Apr 13, 2021 · 84e6900 · 84e6900
1 parent b60824f
commit 84e6900
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pytorch_lightning/plugins/training_type/horovod.py b/pytorch_lightning/plugins/training_type/horovod.py
@@ -136,7 +136,9 @@ def reduce(self, output, group: Optional[Any] = None, reduce_op: Optional[Union[
                 "Unset `group`."
             )
 
-        if reduce_op is None or reduce_op == "sum":
+        if reduce_op in (None, "avg", "mean"):
+            reduce_op = hvd.Average
+        elif reduce_op in ("sum", ReduceOp.SUM):
             reduce_op = hvd.Sum
         elif isinstance(reduce_op, str) and reduce_op in ("avg", "mean"):
             reduce_op = hvd.Average