From f39ee6cf3be349ba4b239c065b93e1e6398096f6 Mon Sep 17 00:00:00 2001 From: Shivam Mehta Date: Tue, 5 Dec 2023 12:10:52 +0000 Subject: [PATCH] Changing while to for for more readibility --- matcha/models/components/flow_matching.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/matcha/models/components/flow_matching.py b/matcha/models/components/flow_matching.py index 4d77547..5cad743 100644 --- a/matcha/models/components/flow_matching.py +++ b/matcha/models/components/flow_matching.py @@ -73,16 +73,14 @@ class BASECFM(torch.nn.Module, ABC): # Or in future might add like a return_all_steps flag sol = [] - steps = 1 - while steps <= len(t_span) - 1: + for step in range(1, len(t_span)): dphi_dt = self.estimator(x, mask, mu, t, spks, cond) x = x + dt * dphi_dt t = t + dt sol.append(x) - if steps < len(t_span) - 1: - dt = t_span[steps + 1] - t - steps += 1 + if step < len(t_span) - 1: + dt = t_span[step + 1] - t return sol[-1]