From e868fca56215bd857dbd65fdc99dee830b264e4e Mon Sep 17 00:00:00 2001 From: Asaf Agami Date: Mon, 13 Apr 2026 18:42:05 +0300 Subject: [PATCH] fix custom_flowmatch_sampler (#783) --- toolkit/samplers/custom_flowmatch_sampler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/toolkit/samplers/custom_flowmatch_sampler.py b/toolkit/samplers/custom_flowmatch_sampler.py index bac7f3fd..fe2e6fa7 100644 --- 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()