Skip to content

Commit

Permalink
Ensure all tensors are on CPU in HitAndRunPolytopeSampler (#2502)
Browse files Browse the repository at this point in the history
Summary:
Addresses #2500

Pull Request resolved: #2502

Reviewed By: SebastianAment

Differential Revision: D62083375

Pulled By: Balandat

fbshipit-source-id: 069f7ff29e9830b0724e742165dd6b90ec41877f
  • Loading branch information
Balandat authored and facebook-github-bot committed Sep 2, 2024
1 parent 1f6bf87 commit 49f7c08
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion botorch/utils/sampling.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 49f7c08

Please sign in to comment.