Add SAAS MTGP to the storage registry (facebook#1619)

Summary: Pull Request resolved: facebook#1619 Reviewed By: dme65 Differential Revision: D45717009 fbshipit-source-id: 94be947003ea72cf268ed9292c872025c14e8eb5
LiamSwayne · May 10, 2023 · 85fb662 · 85fb662
1 parent a7c86cb
commit 85fb662
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/ax/storage/botorch_modular_registry.py b/ax/storage/botorch_modular_registry.py
@@ -31,6 +31,7 @@
     qNoisyExpectedHypervolumeImprovement,
 )
 from botorch.models import SaasFullyBayesianSingleTaskGP
+from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
 
 # BoTorch `Model` imports
 from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP
@@ -85,6 +86,7 @@
     SingleTaskGP: "SingleTaskGP",
     SingleTaskMultiFidelityGP: "SingleTaskMultiFidelityGP",
     SaasFullyBayesianSingleTaskGP: "SaasFullyBayesianSingleTaskGP",
+    SaasFullyBayesianMultiTaskGP: "SaasFullyBayesianMultiTaskGP",
 }