Fix train_Yvar reshaping in SAASBO (#1183)

Summary:
Pull Request resolved: #1183

The fixed noise should have shape `(b) x N` but was passed in as `N x 1`.

Reviewed By: Balandat

Differential Revision: D35702904

fbshipit-source-id: 5ee1887f36a0e83971e1085d8e9ac9175e0b0d36

Author: dme65

botorch/models/fully_bayesian.py

@@ -228,7 +228,11 @@ def load_mcmc_samples(
         ).to(**tkwargs)
         if self.train_Yvar is not None:
             likelihood = FixedNoiseGaussianLikelihood(
-                noise=self.train_Yvar, batch_shape=batch_shape
+                # Reshape to shape `num_mcmc_samples x N`
+                noise=self.train_Yvar.squeeze(-1).expand(
+                    num_mcmc_samples, len(self.train_Yvar)
+                ),
+                batch_shape=batch_shape,
             ).to(**tkwargs)
         else:
             likelihood = GaussianLikelihood(

test/models/test_fully_bayesian.py

@@ -200,10 +200,12 @@ def test_fit_model(self):
             if infer_noise:
                 self.assertEqual(model.likelihood.noise.shape, torch.Size([3, 1]))
             else:
-                self.assertEqual(model.likelihood.noise.shape, torch.Size([n, 1]))
+                self.assertEqual(model.likelihood.noise.shape, torch.Size([3, n]))
                 self.assertTrue(
                     torch.allclose(
-                        train_Yvar.clamp(MIN_INFERRED_NOISE_LEVEL),
+                        train_Yvar.clamp(MIN_INFERRED_NOISE_LEVEL)
+                        .squeeze(-1)
+                        .repeat(3, 1),
                         model.likelihood.noise,
                     )
                 )
@@ -460,7 +462,9 @@ def test_load_samples(self):
                 self.assertTrue(
                     torch.allclose(
                         model.likelihood.noise_covar.noise,
-                        train_Yvar.clamp(MIN_INFERRED_NOISE_LEVEL),
+                        train_Yvar.clamp(MIN_INFERRED_NOISE_LEVEL)
+                        .squeeze(-1)
+                        .repeat(3, 1),
                     )
                 )