Fix None gradient to enable maximum-likelihood training. This fixes issue #59.

flow_matching/solver/ode_solver.py

@@ -159,20 +159,26 @@ def dynamics_func(t, states):
                     # Compute exact divergence
                     div = 0
                     for i in range(ut.flatten(1).shape[1]):
-                        div += gradient(ut[:, i], xt, create_graph=True)[:, i].detach()
+                        g = gradient(ut[:, i], xt, create_graph=True)[:, i]
+                        if not enable_grad:
+                            g = g.detach()
+                        div += g
                 else:
                     # Compute Hutchinson divergence estimator E[z^T D_x(ut) z]
                     ut_dot_z = torch.einsum(
                         "ij,ij->i", ut.flatten(start_dim=1), z.flatten(start_dim=1)
                     )
-                    grad_ut_dot_z = gradient(ut_dot_z, xt)
+                    grad_ut_dot_z = gradient(ut_dot_z, xt, create_graph=enable_grad)
                     div = torch.einsum(
                         "ij,ij->i",
                         grad_ut_dot_z.flatten(start_dim=1),
                         z.flatten(start_dim=1),
                     )
 
-            return ut.detach(), div.detach()
+            if not enable_grad:
+                ut = ut.detach()
+                div = div.detach()
+            return ut, div
 
         y_init = (x_1, torch.zeros(x_1.shape[0], device=x_1.device))
         ode_opts = {"step_size": step_size} if step_size is not None else {}