Make uncertainty toolbox optional

LLNL · Nov 22, 2024 · 47f171a · 47f171a
1 parent ca7bf07
commit 47f171a
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/systems2atoms/hyperplanetree/hyperplanetree/uq.py b/systems2atoms/hyperplanetree/hyperplanetree/uq.py
@@ -8,18 +8,18 @@
 
 import torch
 
-try:
-    from uncertainty_toolbox import get_all_metrics
-except ImportError:
-    raise ImportError('For uncertainty quantification, please install Uncertainty Toolbox: https://uncertainty-toolbox.github.io/')
-
 def calculate_uncertainty_metrics(
         y_pred,
         uncertainty,
         y_true,
         num_bins = 100,
         resolution = 99,
     ):
+
+    try:
+        from uncertainty_toolbox import get_all_metrics
+    except ImportError:
+        raise ImportError('For uncertainty quantification, please install Uncertainty Toolbox: https://uncertainty-toolbox.github.io/')
 
     y_pred = y_pred.flatten()
     uncertainty = uncertainty.flatten()