diff --git a/flow_matching/solver/riemannian_ode_solver.py b/flow_matching/solver/riemannian_ode_solver.py
index 6eb3e5e..6f5ebcb 100644
--- a/flow_matching/solver/riemannian_ode_solver.py
+++ b/flow_matching/solver/riemannian_ode_solver.py
@@ -31,7 +31,6 @@ def __init__(self, manifold: Manifold, velocity_model: ModelWrapper):
         self.manifold = manifold
         self.velocity_model = velocity_model
 
-    @torch.no_grad()
     def sample(
         self,
         x_init: Tensor,