From b61b6d02c017e5c03f5659b491f075a0a13b7ecf Mon Sep 17 00:00:00 2001
From: Marton Havasi <marton@meta.com>
Date: Tue, 17 Dec 2024 10:59:02 +0000
Subject: [PATCH] test lkelihood

---
 tests/solver/test_ode_solver.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/solver/test_ode_solver.py b/tests/solver/test_ode_solver.py
index 394d501..fdee5cc 100644
--- a/tests/solver/test_ode_solver.py
+++ b/tests/solver/test_ode_solver.py
@@ -114,7 +114,7 @@ def dummy_log_p(x: Tensor) -> Tensor:
         self.assertEqual(x_1.shape[0], log_likelihood.shape[0])
 
     def test_compute_likelihood_exact_divergence(self):
-        x_1 = torch.tensor([[0.0, 0.0]])
+        x_1 = torch.tensor([[0.0, 0.0]], requires_grad=True)
         step_size = 0.1
 
         # Define a dummy log probability function
@@ -135,6 +135,10 @@ def dummy_log_p(x: Tensor) -> Tensor:
         self.assertTrue(
             torch.allclose(x_1 - 1.0, x_0, atol=1e-2),
         )
+        log_likelihood.backward()
+        self.assertTrue(
+            torch.allclose(x_1.grad, torch.tensor([1.0, 1.0]), atol=1e-2),
+        )
 
 
 if __name__ == "__main__":