Be more lenient in feature importance plotting for MBM (#1874)

Summary:
Pull Request resolved: #1874

The requirements for BoTorchModel.feature_importances were too strict and causing needless failure to produce plots. Rather than checking that we have a surrogate with a specific key, we just check that there is exactly one surrogate.

Reviewed By: saitcakmak

Differential Revision: D49656886

fbshipit-source-id: 24782696473d186b74bcd83a195754c4d5db3461

ax/models/torch/botorch_modular/model.py

@@ -682,10 +682,10 @@ def feature_importances(self) -> np.ndarray:
             The feature importances as a numpy array of size len(metrics) x 1 x dim
             where each row sums to 1.
         """
-        if list(self.surrogates.keys()) != [Keys.ONLY_SURROGATE]:
+        if len(self.surrogates) != 1:
             raise NotImplementedError("Only support a single surrogate model for now")
-        surrogate = self.surrogates[Keys.ONLY_SURROGATE]
-        return get_feature_importances_from_botorch_model(model=surrogate.model)
+
+        return get_feature_importances_from_botorch_model(model=self.surrogate.model)
 
     @property
     def search_space_digest(self) -> SearchSpaceDigest: