diff --git a/botorch/models/gpytorch.py b/botorch/models/gpytorch.py index 414e7ebe64..0c331c1de7 100644 --- a/botorch/models/gpytorch.py +++ b/botorch/models/gpytorch.py @@ -579,7 +579,7 @@ def posterior( with gpt_posterior_settings(): # only compute what's necessary if output_indices is not None: - mvns = [self.forward_i(i, transformed_X[i]) for i in output_indices] + mvns = [self.models[i](transformed_X[i]) for i in output_indices] if observation_noise is not False: if torch.is_tensor(observation_noise): lh_kwargs = [ diff --git a/test/models/test_gpytorch.py b/test/models/test_gpytorch.py index 12660fd130..6ac7bd9189 100644 --- a/test/models/test_gpytorch.py +++ b/test/models/test_gpytorch.py @@ -458,9 +458,25 @@ def test_model_list_gpytorch_model(self): self.assertIsInstance(posterior, GPyTorchPosterior) self.assertEqual(posterior.mean.shape, torch.Size([2, 2])) # test output indices - posterior = model.posterior(test_X, output_indices=[0]) - self.assertIsInstance(posterior, GPyTorchPosterior) - self.assertEqual(posterior.mean.shape, torch.Size([2, 1])) + for output_indices in ([0], [1], [0, 1]): + posterior_subset = model.posterior( + test_X, output_indices=output_indices + ) + self.assertIsInstance(posterior_subset, GPyTorchPosterior) + self.assertEqual( + posterior_subset.mean.shape, torch.Size([2, len(output_indices)]) + ) + self.assertTrue( + torch.equal( + posterior_subset.mean, posterior.mean[..., output_indices] + ) + ) + self.assertTrue( + torch.equal( + posterior_subset.variance, + posterior.variance[..., output_indices], + ) + ) # test observation noise posterior = model.posterior(test_X, observation_noise=True) self.assertIsInstance(posterior, GPyTorchPosterior)