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()