Merge 992b5b2 into d81a674

botorch/acquisition/penalized.py

@@ -366,6 +366,12 @@ def forward(self, samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
         if self.expand_dim is not None:
             # reshape penalty_obj to match the dim
             penalty_obj = penalty_obj.unsqueeze(self.expand_dim)
+        # this happens when samples is a `q x m`-dim tensor and X is a `q x d`-dim
+        # tensor; obj returned from GenericMCObjective is a `q`-dim tensor and
+        # penalty_obj is a `1 x q`-dim tensor.
+        if obj.ndim == 1:
+            assert penalty_obj.shape == torch.Size([1, samples.shape[-2]])
+            penalty_obj = penalty_obj.squeeze(dim=0)
         return obj - self.regularization_parameter * penalty_obj
 
 

test/acquisition/test_penalized.py

@@ -291,7 +291,7 @@ def test_penalized_mc_objective(self):
             samples = torch.randn(4, 3, device=self.device, dtype=dtype)
             X = torch.randn(4, 5, device=self.device, dtype=dtype)
             penalized_obj = generic_obj(samples) - 0.1 * l1_penalty_obj(X)
-            self.assertTrue(torch.equal(obj(samples, X), penalized_obj))
+            self.assertTrue(torch.equal(obj(samples, X), penalized_obj.squeeze(0)))
             # test 'q x d' Tensor X
             samples = torch.randn(4, 2, 3, device=self.device, dtype=dtype)
             X = torch.randn(2, 5, device=self.device, dtype=dtype)