Fix broken cuda tests (#1115)

Summary: Pull Request resolved: #1115 See title Reviewed By: saitcakmak Differential Revision: D34761452 fbshipit-source-id: b29f6d86d3aaa8b8b16e6c660d8d8d7dac65758d
pytorch · Mar 9, 2022 · 18a3b26 · 18a3b26
1 parent 21b4616
commit 18a3b26
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 10 deletions.
diff --git a/botorch/models/multitask.py b/botorch/models/multitask.py
@@ -473,7 +473,7 @@ def __init__(
         if task_covar_prior is None:
             task_covar_prior = LKJCovariancePrior(
                 n=num_tasks,
-                eta=kwargs.get("eta", 1.5),
+                eta=torch.tensor(kwargs.get("eta", 1.5)).to(train_X),
                 sd_prior=kwargs.get(
                     "sd_prior",
                     SmoothedBoxPrior(math.exp(-6), math.exp(1.25), 0.05),

diff --git a/test/acquisition/multi_objective/test_utils.py b/test/acquisition/multi_objective/test_utils.py
@@ -129,7 +129,9 @@ def test_prune_inferior_points_multi_objective(self):
                 )
             if self.device.type == "cuda":
                 # sorting has different order on cuda
-                self.assertTrue(torch.equal(X_pruned, torch.stack([X[2], X[1]], dim=0)))
+                self.assertTrue(
+                    torch.equal(X_pruned, X[[2, 1]]) or torch.equal(X_pruned, X[[1, 2]])
+                )
             else:
                 self.assertTrue(torch.equal(X_pruned, X[:2]))
             # test that zero-probability is in fact pruned

diff --git a/test/acquisition/test_max_value_entropy_search.py b/test/acquisition/test_max_value_entropy_search.py
@@ -158,7 +158,9 @@ def test_q_max_value_entropy(self):
 
             # Test with multi-output model w/ transform.
             mm = MESMockModel(num_outputs=2)
-            pt = ScalarizedPosteriorTransform(weights=torch.ones(2, dtype=dtype))
+            pt = ScalarizedPosteriorTransform(
+                weights=torch.ones(2, device=self.device, dtype=dtype)
+            )
             for gumbel in (True, False):
                 qMVE = qMaxValueEntropy(
                     mm,
@@ -225,7 +227,9 @@ def test_q_lower_bound_max_value_entropy(self):
 
             # Test with multi-output model w/ transform.
             mm = MESMockModel(num_outputs=2)
-            pt = ScalarizedPosteriorTransform(weights=torch.ones(2, dtype=dtype))
+            pt = ScalarizedPosteriorTransform(
+                weights=torch.ones(2, device=self.device, dtype=dtype)
+            )
             qGIBBON = qLowerBoundMaxValueEntropy(
                 mm,
                 candidate_set,
@@ -270,7 +274,9 @@ def test_q_multi_fidelity_max_value_entropy(self):
 
             # Test with multi-output model w/ transform.
             mm = MESMockModel(num_outputs=2)
-            pt = ScalarizedPosteriorTransform(weights=torch.ones(2, dtype=dtype))
+            pt = ScalarizedPosteriorTransform(
+                weights=torch.ones(2, device=self.device, dtype=dtype)
+            )
             qMF_MVE = qMultiFidelityMaxValueEntropy(
                 model=mm,
                 candidate_set=candidate_set,
@@ -284,13 +290,15 @@ def test_sample_max_value_Gumbel(self):
         for dtype in (torch.float, torch.double):
             torch.manual_seed(7)
             mm = MESMockModel()
-            candidate_set = torch.rand(3, 10, 2, dtype=dtype)
+            candidate_set = torch.rand(3, 10, 2, device=self.device, dtype=dtype)
             samples = _sample_max_value_Gumbel(mm, candidate_set, 5)
             self.assertEqual(samples.shape, torch.Size([5, 3]))
 
             # Test with multi-output model w/ transform.
             mm = MESMockModel(num_outputs=2)
-            pt = ScalarizedPosteriorTransform(weights=torch.ones(2, dtype=dtype))
+            pt = ScalarizedPosteriorTransform(
+                weights=torch.ones(2, device=self.device, dtype=dtype)
+            )
             samples = _sample_max_value_Gumbel(
                 mm, candidate_set, 5, posterior_transform=pt
             )
@@ -300,13 +308,15 @@ def test_sample_max_value_Thompson(self):
         for dtype in (torch.float, torch.double):
             torch.manual_seed(7)
             mm = MESMockModel()
-            candidate_set = torch.rand(3, 10, 2, dtype=dtype)
+            candidate_set = torch.rand(3, 10, 2, device=self.device, dtype=dtype)
             samples = _sample_max_value_Thompson(mm, candidate_set, 5)
             self.assertEqual(samples.shape, torch.Size([5, 3]))
 
             # Test with multi-output model w/ transform.
             mm = MESMockModel(num_outputs=2)
-            pt = ScalarizedPosteriorTransform(weights=torch.ones(2, dtype=dtype))
+            pt = ScalarizedPosteriorTransform(
+                weights=torch.ones(2, device=self.device, dtype=dtype)
+            )
             samples = _sample_max_value_Thompson(
                 mm, candidate_set, 5, posterior_transform=pt
             )

diff --git a/test/models/test_multitask.py b/test/models/test_multitask.py
@@ -706,7 +706,7 @@ def test_KroneckerMultiTaskGP_custom(self):
             )
             task_covar_prior = LKJCovariancePrior(
                 n=2,
-                eta=0.5,
+                eta=torch.tensor(0.5, **tkwargs),
                 sd_prior=SmoothedBoxPrior(math.exp(-3), math.exp(2), 0.1),
             )
             model_kwargs = {