Actually use the correct timestep sampling instead of calculating it and moving on lol. Tested a few with it and it seems to work better.

toolkit/samplers/custom_flowmatch_sampler.py

@@ -55,4 +55,6 @@ class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
         # Sort the timesteps in descending order
         timesteps, _ = torch.sort(timesteps, descending=True)
 
+        self.timesteps = timesteps.to(device=device)
+
         return timesteps