move tensors to correct device (#1)

prescient-design · Jul 9, 2024 · 4d94731 · 4d94731
1 parent 464216c
commit 4d94731
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/holo/test_functions/closed_form/_rough_mt_fuji.py b/holo/test_functions/closed_form/_rough_mt_fuji.py
@@ -60,15 +60,18 @@ def _optimal_value(self):
     def optimal_solution(self):
         soln = self.centroids.clone()
         mask = self._random_term - self._additive_term > 0
+        mask = mask.to(device=soln.device)
         soln = torch.where(mask, torch.ones_like(soln), soln)
         mask = self._random_term + self._additive_term < 0
+        mask = mask.to(device=soln.device)
         soln = torch.where(mask, torch.zeros_like(soln), soln)
 
         return soln
 
     def to(self, device, dtype):
         self.centroids = self.centroids.to(device, dtype)
         self._generator = torch.Generator(device=device).manual_seed(self._random_seed)
+        self._random_term = self._random_term.to(device)
         return self
 
     def __repr__(self):