Fix test failure in test_evaluate_qMFKG (#618)

Summary:
Caused by a PR race condition between #594 and #588

Pull Request resolved: #618

Reviewed By: qingfeng10

Differential Revision: D25386475

Pulled By: Balandat

fbshipit-source-id: 95f445c5a03e07b2c35eab7134b43b39ec25c472

Author: Balandat

botorch/utils/testing.py

@@ -169,6 +169,11 @@ def num_outputs(self) -> int:
         event_shape = self._posterior.event_shape
         return event_shape[-1] if len(event_shape) > 0 else 0
 
+    @property
+    def batch_shape(self) -> torch.Size:
+        event_shape = self._posterior.event_shape
+        return event_shape[:-2]
+
     def state_dict(self) -> None:
         pass
 

test/acquisition/test_knowledge_gradient.py

@@ -406,13 +406,11 @@ def test_evaluate_qMFKG(self):
         for dtype in (torch.float, torch.double):
             mean = torch.zeros(1, 1, 1, device=self.device, dtype=dtype)
             mm = MockModel(MockPosterior(mean=mean))
-            mm._input_batch_shape = torch.Size([1])
             cau = GenericCostAwareUtility(mock_util)
             n_f = 4
             mean = torch.rand(n_f, 2, 1, 1, device=self.device, dtype=dtype)
             variance = torch.rand(n_f, 2, 1, 1, device=self.device, dtype=dtype)
             mfm = MockModel(MockPosterior(mean=mean, variance=variance))
-            mfm._input_batch_shape = torch.Size([n_f, 2])
             with ExitStack() as es:
                 patch_f = es.enter_context(
                     mock.patch.object(MockModel, "fantasize", return_value=mfm)