Update gpytorch/test/base_keops_test_case.py

Co-authored-by: Max Balandat <Balandat@users.noreply.github.com>
cornellius-gp · Jan 28, 2025 · c39666c · c39666c
1 parent fe4619b
commit c39666c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/gpytorch/test/base_keops_test_case.py b/gpytorch/test/base_keops_test_case.py
@@ -46,6 +46,7 @@ def test_forward_x1_eq_x2(self, ard=False, use_keops=True, **kwargs):
                 d1 = kern1(x1, x1).diagonal(dim1=-1, dim2=-2)
                 d2 = kern2(x1, x1).diagonal(dim1=-1, dim2=-2)
                 self.assertLess(torch.norm(d1 - d2), 1e-4)
+                self.assertTrue(torch.equal(k1.diag(), d1))
 
         if use_keops:
             self.assertTrue(keops_mock.called)