Extract scalar only when necessary

Author: AdrianSosic

botorch/optim/optimize.py

@@ -196,7 +196,9 @@ 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[0]
+    if acq_value.ndim == 1:
+        acq_value = acq_value[0]
+    return X, acq_value
 
 
 def _validate_sequential_inputs(opt_inputs: OptimizeAcqfInputs) -> None: