From f28558a4c345bf0b6b36d94c35534f38fbf8f335 Mon Sep 17 00:00:00 2001
From: timonpalm <timon.palm@gmx.de>
Date: Thu, 31 Jul 2025 15:36:31 +0200
Subject: [PATCH] t on right device

---
 flow_matching/solver/ode_solver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flow_matching/solver/ode_solver.py b/flow_matching/solver/ode_solver.py
index 8997506..1c2bbd8 100644
--- a/flow_matching/solver/ode_solver.py
+++ b/flow_matching/solver/ode_solver.py
@@ -181,7 +181,7 @@ def dynamics_func(t, states):
             sol, log_det = odeint(
                 dynamics_func,
                 y_init,
-                time_grid,
+                time_grid.to(x_1.device),
                 method=method,
                 options=ode_opts,
                 atol=atol,