From 6ff3789e151c7e49995576a4a39399342b856618 Mon Sep 17 00:00:00 2001 From: Kohya S Date: Tue, 22 Aug 2023 23:37:54 +0900 Subject: [PATCH] fix steps wasn't reset --- node_control_net_lllite.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/node_control_net_lllite.py b/node_control_net_lllite.py index cc33b95..cb67f7d 100644 --- a/node_control_net_lllite.py +++ b/node_control_net_lllite.py @@ -210,9 +210,11 @@ class LLLiteModule(torch.nn.Module): if self.is_first and self.current_step == self.end_step: print(f"end LLLite: step {self.current_step}") self.current_step += 1 + if self.current_step >= self.num_steps: + self.current_step = 0 # reset return torch.zeros_like(x) else: - if self.is_first and self.current_step ==self.start_step: + if self.is_first and self.current_step == self.start_step: print(f"start LLLite: step {self.current_step}") self.current_step += 1 if self.current_step >= self.num_steps: