fix

CNChTu · Oct 7, 2024 · b93e925 · b93e925
1 parent 3944dcf
commit b93e925
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/diffusion/reflow/reflow.py b/diffusion/reflow/reflow.py
@@ -55,7 +55,7 @@ def reflow_consistency_loss(self, x_1, t_a, t_b, cond, loss_type=None):
         x_t_a = x_0 + t_a[:, None, None, None] * (x_1 - x_0)
         x_t_b = x_0 + t_b[:, None, None, None] * (x_1 - x_0)
         v_pred_a = self.velocity_fn(x_t_a, 1000 * t_a, cond)
-        v_pred_b = self.velocity_fn(x_t_b, 1000 * t_b, cond)
+        v_pred_b = self.velocity_fn(x_t_b, 1000 * t_b, cond).detach()
 
         if loss_type is None:
             loss_type = self.loss_type