Merge branch 'master' into fix

Author: Balandat

gpytorch/lazy/cat_lazy_tensor.py

@@ -181,7 +181,10 @@ def _get_indices(self, row_index, col_index, *batch_indices):
         if len(res_list) == 1:
             return res_list[0].view(target_shape).to(self.device)
         else:
-            return torch.cat(res_list).view(target_shape).to(self.device)
+            # Explicitly move tensors to one device as torch.cat no longer moves tensors:
+            # https://github.com/pytorch/pytorch/issues/35045
+            res_list = [lazy_tensor.to(self.device) for lazy_tensor in res_list]
+            return torch.cat(res_list).view(target_shape)
 
     def _getitem(self, row_index, col_index, *batch_indices):
         indices = [*batch_indices, row_index, col_index]

gpytorch/utils/permutation.py

@@ -76,7 +76,8 @@ def apply_permutation(
         right_permutation = torch.arange(matrix.size(-1), device=matrix.device)
 
     # Apply permutations
-    return delazify(matrix.__getitem__((*batch_idx, left_permutation.unsqueeze(-1), right_permutation.unsqueeze(-2))))
+    res = delazify(matrix.__getitem__((*batch_idx, left_permutation.unsqueeze(-1), right_permutation.unsqueeze(-2))))
+    return res.to(device=matrix.device)
 
 
 def inverse_permutation(permutation):