Minor tweak to MARS.set_baseline_Y to reduce memory overhead (#1346)

Summary:
Pull Request resolved: #1346

Unsqueeze before evaluating `model.posterior` in `MARS.set_baseline_Y` to avoid creating large matrices. This helps reduce the memory overhead with large models and helps avoid OOMs.

Reviewed By: Balandat

Differential Revision: D38648419

fbshipit-source-id: c4bd5646e53625f640f545b9a33979d8458c885a

Author: saitcakmak

botorch/acquisition/multi_objective/multi_output_risk_measures.py

@@ -505,7 +505,7 @@ def forward(self, samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
             if self.expectation:
                 padded_mvar_list.append(mvar_.mean(dim=0))
             else:
-                # Repeat the last entry to make `mvar_set` `n_w x m`.
+                # Repeat the last entry to make `mvar_set` `pad_size x m`.
                 repeats_needed = pad_size - mvar_.shape[0]
                 padded_mvar_list.append(
                     torch.cat([mvar_, mvar_[-1].expand(repeats_needed, m)], dim=0)
@@ -578,7 +578,7 @@ def set_baseline_Y(self, model: Model, X_baseline: Tensor) -> None:
             X_baseline: An `n x d`-dim tensor of previously evaluated points.
         """
         with torch.no_grad():
-            Y = model.posterior(X_baseline).mean
+            Y = model.posterior(X_baseline.unsqueeze(-2)).mean.squeeze(-2)
         Y = self.preprocessing_function(Y)
         Y = self.mvar(Y).view(-1, Y.shape[-1])
         Y = Y[is_non_dominated(Y)]