From fbed8568fbefc61484ec4ad3b4aa907a4800238b Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 11 Aug 2024 11:10:37 -0600 Subject: [PATCH] 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 | 2 ++ 1 file changed, 2 insertions(+) diff --git a/toolkit/samplers/custom_flowmatch_sampler.py b/toolkit/samplers/custom_flowmatch_sampler.py index a38fcce0..0a1d7f45 100644 --- a/toolkit/samplers/custom_flowmatch_sampler.py +++ b/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