Merge pull request #1712 from wjmaddox/cat_rows_fix

Add device settings to cat_rows

Author: Balandat

gpytorch/lazy/lazy_tensor.py

@@ -780,9 +780,9 @@ def cat_rows(self, cross_mat, new_mat, generate_roots=True, **root_decomp_kwargs
             A = self
 
         # form matrix C = [A B; B^T D], where A = self, B = cross_mat, D = new_mat
-        upper_row = CatLazyTensor(A, B, dim=-2)
-        lower_row = CatLazyTensor(B.transpose(-1, -2), D, dim=-2)
-        new_lazy_tensor = CatLazyTensor(upper_row, lower_row, dim=-1)
+        upper_row = CatLazyTensor(A, B, dim=-2, output_device=A.device)
+        lower_row = CatLazyTensor(B.transpose(-1, -2), D, dim=-2, output_device=A.device)
+        new_lazy_tensor = CatLazyTensor(upper_row, lower_row, dim=-1, output_device=A.device)
 
         # if the old lazy tensor does not have either a root decomposition or a root inverse decomposition
         # don't create one