Fixed a bug when loading horovod

RandomDefaultUser · Oct 12, 2023 · 159f2c5 · 159f2c5
1 parent b697f12
commit 159f2c5
Showing 1 changed file with 20 additions and 12 deletions.
diff --git a/mala/common/parameters.py b/mala/common/parameters.py
@@ -6,8 +6,10 @@
 import pickle
 from time import sleep
 
+horovod_available = False
 try:
     import horovod.torch as hvd
+    horovod_available = True
 except ModuleNotFoundError:
     pass
 import numpy as np
@@ -1257,19 +1259,25 @@ def use_horovod(self):
 
     @use_horovod.setter
     def use_horovod(self, value):
-        if value:
-            hvd.init()
+        if value is False:
+            self._use_horovod = False
+        else:
+            if horovod_available:
+                hvd.init()
+                # Invalidate, will be updated in setter.
+                set_horovod_status(value)
+                self.device = None
+                self._use_horovod = value
+                self.network._update_horovod(self.use_horovod)
+                self.descriptors._update_horovod(self.use_horovod)
+                self.targets._update_horovod(self.use_horovod)
+                self.data._update_horovod(self.use_horovod)
+                self.running._update_horovod(self.use_horovod)
+                self.hyperparameters._update_horovod(self.use_horovod)
+            else:
+                parallel_warn("Horovod requested, but not installed found. "
+                              "MALA will operate without horovod only.")
 
-        # Invalidate, will be updated in setter.
-        set_horovod_status(value)
-        self.device = None
-        self._use_horovod = value
-        self.network._update_horovod(self.use_horovod)
-        self.descriptors._update_horovod(self.use_horovod)
-        self.targets._update_horovod(self.use_horovod)
-        self.data._update_horovod(self.use_horovod)
-        self.running._update_horovod(self.use_horovod)
-        self.hyperparameters._update_horovod(self.use_horovod)
 
     @property
     def device(self):