Merge pull request #30 from wassimseif/scheduler_fix

Fixed a bug where the scheduler would get None
lucidrains · Oct 1, 2023 · 4f04a03 · 4f04a03
2 parents f748f5d + 59dbc3b
commit 4f04a03
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/.gitignore b/.gitignore
@@ -157,4 +157,4 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/
diff --git a/voicebox_pytorch/trainer.py b/voicebox_pytorch/trainer.py
@@ -105,7 +105,7 @@ def __init__(
 
         self.lr = lr
         self.initial_lr = initial_lr
-        self.scheduler = CosineAnnealingLR(self.optim, T_max = num_train_steps)
+
 
         # max grad norm
 
@@ -135,7 +135,7 @@ def __init__(
             self.num_train_steps = len(dataset) // batch_size * num_epochs
         else:
             self.num_train_steps = num_train_steps
-
+        self.scheduler = CosineAnnealingLR(self.optim, T_max=self.num_train_steps)
         self.num_warmup_steps = num_warmup_steps if exists(num_warmup_steps) else 0
 
         # dataloader