fix uint8 bug in CV (#224)

Summary:
Pull Request resolved: #224

see title

Reviewed By: Balandat

Differential Revision: D16721744

fbshipit-source-id: dab8d9c9d1b85a87367e6ddd2fe4c68f43c06583

Author: sdaulton

botorch/cross_validation.py

@@ -68,6 +68,7 @@ def gen_loo_cv_folds(
         >>> cv_folds = gen_loo_cv_folds(train_X, train_Y)
     """
     masks = torch.eye(train_X.shape[-2], dtype=torch.uint8, device=train_X.device)
+    masks = masks.to(dtype=torch.bool)
     if train_Y.dim() < train_X.dim():
         # add output dimension
         train_Y = train_Y.unsqueeze(-1)

test/test_cross_validation.py

@@ -43,6 +43,35 @@ def test_single_task_batch_cv(self, cuda=False):
                     noiseless_cv_folds = gen_loo_cv_folds(
                         train_X=train_X, train_Y=train_Y
                     )
+                    # check shapes
+                    expected_shape_train_X = batch_shape + torch.Size(
+                        [n, n - 1, train_X.shape[-1]]
+                    )
+                    expected_shape_test_X = batch_shape + torch.Size(
+                        [n, 1, train_X.shape[-1]]
+                    )
+                    self.assertEqual(
+                        noiseless_cv_folds.train_X.shape, expected_shape_train_X
+                    )
+                    self.assertEqual(
+                        noiseless_cv_folds.test_X.shape, expected_shape_test_X
+                    )
+
+                    expected_shape_train_Y = batch_shape + torch.Size(
+                        [n, n - 1, num_outputs]
+                    )
+                    expected_shape_test_Y = batch_shape + torch.Size(
+                        [n, 1, num_outputs]
+                    )
+
+                    self.assertEqual(
+                        noiseless_cv_folds.train_Y.shape, expected_shape_train_Y
+                    )
+                    self.assertEqual(
+                        noiseless_cv_folds.test_Y.shape, expected_shape_test_Y
+                    )
+                    self.assertIsNone(noiseless_cv_folds.train_Yvar)
+                    self.assertIsNone(noiseless_cv_folds.test_Yvar)
                     # Test SingleTaskGP
                     with warnings.catch_warnings():
                         warnings.filterwarnings("ignore", category=OptimizationWarning)
@@ -60,6 +89,21 @@ def test_single_task_batch_cv(self, cuda=False):
                     noisy_cv_folds = gen_loo_cv_folds(
                         train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar
                     )
+                    # check shapes
+                    self.assertEqual(
+                        noisy_cv_folds.train_X.shape, expected_shape_train_X
+                    )
+                    self.assertEqual(noisy_cv_folds.test_X.shape, expected_shape_test_X)
+                    self.assertEqual(
+                        noisy_cv_folds.train_Y.shape, expected_shape_train_Y
+                    )
+                    self.assertEqual(noisy_cv_folds.test_Y.shape, expected_shape_test_Y)
+                    self.assertEqual(
+                        noisy_cv_folds.train_Yvar.shape, expected_shape_train_Y
+                    )
+                    self.assertEqual(
+                        noisy_cv_folds.test_Yvar.shape, expected_shape_test_Y
+                    )
                     with warnings.catch_warnings():
                         warnings.filterwarnings("ignore", category=OptimizationWarning)
                         cv_results = batch_cross_validation(