Fix custom modules import error for predict

Signed-off-by: Benedikt Blumenstiel <[email protected]>
IBM · Dec 4, 2024 · dc6d188 · dc6d188
1 parent 4726ecf
commit dc6d188
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/terratorch/cli_tools.py b/terratorch/cli_tools.py
@@ -389,6 +389,8 @@ def instantiate_classes(self) -> None:
             custom_modules_path = self.config.fit.custom_modules_path
         elif hasattr(self.config, "test") and hasattr(self.config.test, "custom_modules_path"):
             custom_modules_path = self.config.test.custom_modules_path
+        elif hasattr(self.config, "predict") and hasattr(self.config.predict, "custom_modules_path"):
+            custom_modules_path = self.config.predict.custom_modules_path
         else:
             default_path = Path(".") / "custom_modules"
             custom_modules_path = os.environ.get("TERRATORCH_CUSTOM_MODULE_PATH", default_path)