fix: only overwrite raw_samples within each q loop

pytorch · Nov 22, 2024 · 1ee598c · 1ee598c
1 parent 49bd6bf
commit 1ee598c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/botorch/optim/optimize_homotopy.py b/botorch/optim/optimize_homotopy.py
@@ -157,7 +157,6 @@ def optimize_acqf_homotopy(
     """
     shared_optimize_acqf_kwargs = {
         "num_restarts": num_restarts,
-        "raw_samples": raw_samples,
         "inequality_constraints": inequality_constraints,
         "equality_constraints": equality_constraints,
         "nonlinear_inequality_constraints": nonlinear_inequality_constraints,
@@ -178,6 +177,7 @@ def optimize_acqf_homotopy(
 
     for _ in range(q):
         candidates = batch_initial_conditions
+        q_raw_samples = raw_samples
         homotopy.restart()
 
         while not homotopy.should_stop:
@@ -187,13 +187,14 @@ def optimize_acqf_homotopy(
                 q=1,
                 options=options,
                 batch_initial_conditions=candidates,
+                raw_samples=q_raw_samples,
                 **shared_optimize_acqf_kwargs,
             )
             homotopy.step()
 
             # Set raw_samples to None such that pruned restarts are not repopulated
             # at each step in the homotopy.
-            shared_optimize_acqf_kwargs["raw_samples"] = None
+            q_raw_samples = None
 
             # Prune candidates
             candidates = prune_candidates(
@@ -208,6 +209,7 @@ def optimize_acqf_homotopy(
             bounds=bounds,
             q=1,
             options=final_options,
+            raw_samples=q_raw_samples,
             batch_initial_conditions=candidates,
             **shared_optimize_acqf_kwargs,
         )