Merge pull request #68 from francesco-vaselli/main

add padding of t in TargetConditionalFlowMatcher (fix bug)
atong01 · Nov 10, 2023 · 21cd0c8 · 21cd0c8
2 parents a7a6d83 + 5b3a22e
commit 21cd0c8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchcfm/conditional_flow_matching.py b/torchcfm/conditional_flow_matching.py
@@ -326,6 +326,7 @@ def compute_mu_t(self, x0, x1, t):
         [2] Flow Matching for Generative Modelling, ICLR, Lipman et al.
         """
         del x0
+        t = pad_t_like_x(t, x1)
         return t * x1
 
     def compute_sigma_t(self, t):
@@ -369,6 +370,7 @@ def compute_conditional_flow(self, x0, x1, t, xt):
         [1] Flow Matching for Generative Modelling, ICLR, Lipman et al.
         """
         del x0
+        t = pad_t_like_x(t, x1)
         return (x1 - (1 - self.sigma) * xt) / (1 - (1 - self.sigma) * t)