test lkelihood

facebookresearch · Dec 17, 2024 · b61b6d0 · b61b6d0
1 parent be97497
commit b61b6d0
Showing 1 changed file with 5 additions and 1 deletion.
diff --git 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__":