Fix error when using gpu

pytorch · Dec 21, 2023 · e258563 · e258563
1 parent 545e388
commit e258563
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/botorch/acquisition/augmented_multisource.py b/botorch/acquisition/augmented_multisource.py
@@ -79,7 +79,7 @@ def forward(self, X: Tensor) -> Tensor:
             A `(b1 x ... bk)`-dim tensor of Augmented Upper Confidence Bound values at
             the given design points `X`.
         """
-        alpha = torch.zeros(X.shape[0], dtype=X.dtype)
+        alpha = torch.zeros(X.shape[0], dtype=X.dtype, device=X.device)
         agp_mean, agp_sigma = self._mean_and_sigma(X[..., :-1])
         cb = (self.best_f if self.maximize else -self.best_f) + (
             (agp_mean if self.maximize else -agp_mean) + self.beta.sqrt() * agp_sigma

diff --git a/test/models/test_gp_regression_multisource.py b/test/models/test_gp_regression_multisource.py
@@ -33,14 +33,13 @@
 
 
 def _get_random_data_with_source(batch_shape, n, d, n_source, q=1, **tkwargs):
-    dtype = tkwargs.get("dtype", torch.float32)
     rep_shape = batch_shape + torch.Size([1, 1])
     bounds = torch.stack([torch.zeros(d), torch.ones(d)])
     bounds[-1, -1] = n_source - 1
     train_x = (
-        get_random_x_for_agp(n=n, bounds=bounds, q=q).repeat(rep_shape).type(dtype)
+        get_random_x_for_agp(n=n, bounds=bounds, q=q).repeat(rep_shape).to(**tkwargs)
     )
-    train_y = torch.sin(train_x[..., :1] * (2 * math.pi)).type(dtype)
+    train_y = torch.sin(train_x[..., :1] * (2 * math.pi)).to(**tkwargs)
     train_y = train_y + 0.2 * torch.randn(n, 1, **tkwargs).repeat(rep_shape)
     return train_x, train_y