fix steps wasn't reset

2023-08-22 23:37:54 +09:00
parent 159377373c
commit 6ff3789e15
1 changed files with 3 additions and 1 deletions
@@ -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: