Merge pull request #1019 from UaRuairc/patch-1

Fix res_multistep.step to set sigma_next correctly when ending early
2025-08-11 16:24:38 +03:00
parent 057ee642f6 1e091d181b
commit bf84e99dec
1 changed files with 1 additions and 1 deletions
@@ -51,7 +51,7 @@ class FlowMatchSchedulerResMultistep():
        
        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
        sigma_prev = self.sigmas[timestep_id - 1].reshape(-1, 1, 1, 1) if timestep_id > 0 else sigma
-        if (timestep_id + 1 >= len(self.timesteps)).any():
+        if (timestep_id + 1 >= len(self.sigmas)).any():
            sigma_next = torch.tensor(0)
        else:
            sigma_next = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)