fix custom_flowmatch_sampler (#783)

2026-04-13 18:42:05 +03:00 · 2026-04-13 18:42:05 +03:00 · e868fca562
parent 233e292256
commit e868fca562
1 changed files with 1 additions and 1 deletions
--- a/toolkit/samplers/custom_flowmatch_sampler.py
+++ b/toolkit/samplers/custom_flowmatch_sampler.py
@ -68,7 +68,7 @@ class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
                device=timesteps.device,
                dtype=timesteps.dtype
            )
-        if v2:
+        elif v2:
            weights = self.linear_timesteps_weights2[step_indices].flatten()
        else:
            weights = self.linear_timesteps_weights[step_indices].flatten()