Moved cache_root to propery place under noisy aqs

MSDLLCpapers · Sep 17, 2024 · 8831fa3 · 8831fa3
1 parent 17965cd
commit 8831fa3
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/obsidian/optimizer/bayesian.py b/obsidian/optimizer/bayesian.py
@@ -534,7 +534,9 @@ def _parse_aq_kwargs(self,
         # Noisy aqs require X_train reference
         if aq in ['NEI', 'NEHVI', 'NParEGO']:
             aq_kwargs['X_baseline'] = X_baseline
-
+            if any(isinstance(m, EnsembleModel) for m in model.models):
+                aq_kwargs['cache_root'] = False
+
         # Hypervolume requires reference point
         if aq in ['EHVI', 'NEHVI']:
 
@@ -571,8 +573,8 @@ def _parse_aq_kwargs(self,
                 w = w/torch.sum(torch.abs(w))
             aq_kwargs['scalarization_weights'] = w
 
-        if any(isinstance(m, EnsembleModel) for m in model.models):
-            aq_kwargs['cache_root'] = False
+        if aq == 'SF':
+            aq_kwargs['X_baseline'] = X_baseline
 
         return aq_kwargs