From 506c5ba83d53326fd071869666110700ae6add99 Mon Sep 17 00:00:00 2001
From: Maxim Ziatdinov <34245227+ziatdinovmax@users.noreply.github.com>
Date: Sun, 20 Aug 2023 19:32:57 -0400
Subject: [PATCH] Fix bug in UE

After the last update it was actually computing ucb
---
 gpax/acquisition/acquisition.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/gpax/acquisition/acquisition.py b/gpax/acquisition/acquisition.py
index df9280b..73be596 100644
--- a/gpax/acquisition/acquisition.py
+++ b/gpax/acquisition/acquisition.py
@@ -15,8 +15,7 @@
 import numpy as onp
 
 from ..models.gp import ExactGP
-from ..utils import random_sample_dict
-from .base_acq import ei, ucb, poi
+from .base_acq import ei, ucb, poi, ue
 from .penalties import compute_penalty
 
 
@@ -321,7 +320,7 @@ def UE(rng_key: jnp.ndarray,
         warnings.warn("`rng_key` is deprecated and will be removed in future versions. "
                       "It's no longer used.", DeprecationWarning, stacklevel=2)
     return compute_acquisition(
-        model, X, ucb, noiseless,
+        model, X, ue, noiseless,
         penalty=penalty, recent_points=recent_points,
         grid_indices=grid_indices, penalty_factor=penalty_factor,
         **kwargs)