Ensure all tensors are on CPU in HitAndRunPolytopeSampler

Addresses #2500
pytorch · Sep 1, 2024 · d62fd3c · d62fd3c
1 parent 1f6bf87
commit d62fd3c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/botorch/utils/sampling.py b/botorch/utils/sampling.py
@@ -683,7 +683,9 @@ def draw(self, n: int = 1) -> Tensor:
             # Run this on the cpu since there is a lot of looping going on
             A=self.new_A.cpu(),
             b=(self.b - self.A @ self.x0).cpu(),
-            x0=torch.zeros((self.nullC.size(1), 1), dtype=self.A.dtype),
+            x0=torch.zeros(
+                (self.nullC.size(1), 1), dtype=self.A.dtype, device=torch.device("cpu")
+            ),
             n=n,
             n0=self.n_burnin if self.num_samples_generated == 0 else 0,
             n_thinning=self.n_thinning,