Fix Checkpoint issue when using Horovod distributed backend (PyTorchLightning#6947) (#6958)

Co-Authored-By: Adrian Wälchli <aedu.waelchli@gmail.com> Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>
2021-04-13 11:18:52 +02:00 · 2021-04-13 11:18:52 +02:00 · b37b58a73e
parent da1ac3a530
commit b37b58a73e
1 changed files with 1 additions and 1 deletions
--- a/pytorch_lightning/plugins/training_type/horovod.py
+++ b/pytorch_lightning/plugins/training_type/horovod.py
@ -150,7 +150,7 @@ class HorovodPlugin(ParallelPlugin):

        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}")