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__":