Detach fval only once in torch_minimize

pytorch · Sep 12, 2024 · 5e64cb9 · 5e64cb9
1 parent d974542
commit 5e64cb9
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/botorch/optim/core.py b/botorch/optim/core.py
@@ -192,17 +192,17 @@ def torch_minimize(
         else {name: limits for name, limits in bounds.items() if name in parameters}
     )
     for step in range(1, step_limit + 1):
-        fval, _ = closure()
+        fval = closure()[0].detach()
         runtime = monotonic() - start_time
         result = OptimizationResult(
             step=step,
-            fval=fval.detach().cpu().item(),
+            fval=fval.cpu().item(),
             status=OptimizationStatus.RUNNING,
             runtime=runtime,
         )
 
         # TODO: Update stopping_criterion API to return a message.
-        if stopping_criterion and stopping_criterion(fval.detach()):
+        if stopping_criterion and stopping_criterion(fval):
             result.status = OptimizationStatus.STOPPED
             result.message = "`torch_minimize` stopped due to `stopping_criterion`."