Drop output dimension of acquisition function values

pytorch · Feb 12, 2025 · f23e90b · f23e90b
1 parent 8c73830
commit f23e90b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/botorch/optim/optimize.py b/botorch/optim/optimize.py
@@ -192,7 +192,7 @@ def _optimize_acqf_all_features_fixed(
     X = X.expand(q, *X.shape)
     with torch.no_grad():
         acq_value = acq_function(X)
-    return X, acq_value
+    return X, acq_value[0]
 
 
 def _validate_sequential_inputs(opt_inputs: OptimizeAcqfInputs) -> None: