From 47f171a58cbdec432cce030a2085881b495cad79 Mon Sep 17 00:00:00 2001
From: Ethan Sunshine <esunshin@andrew.cmu.edu>
Date: Fri, 22 Nov 2024 13:52:21 +0000
Subject: [PATCH] Make uncertainty toolbox optional

---
 systems2atoms/hyperplanetree/hyperplanetree/uq.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/systems2atoms/hyperplanetree/hyperplanetree/uq.py b/systems2atoms/hyperplanetree/hyperplanetree/uq.py
index 81840c0..68eed9a 100644
--- a/systems2atoms/hyperplanetree/hyperplanetree/uq.py
+++ b/systems2atoms/hyperplanetree/hyperplanetree/uq.py
@@ -8,11 +8,6 @@
 
 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,
@@ -20,6 +15,11 @@ def calculate_uncertainty_metrics(
         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()