github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fixed bug the kept learning rates the same

Browse Source
This commit is contained in:
Jaret Burkett
2024-03-06 09:23:32 -07:00
parent b01e8d889a
commit f1cb87fe9e
2 changed files with 18 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -1363,7 +1363,13 @@ class SDTrainer(BaseSDTrainProcess):
# flush()
if not self.is_grad_accumulation_step:
torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm)
# torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm)
# fix this for multi params
if isinstance(self.params[0], dict):
for i in range(len(self.params)):
torch.nn.utils.clip_grad_norm_(self.params[i]['params'], self.train_config.max_grad_norm)
else:
torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm)
# only step if we are not accumulating
with self.timer('optimizer_step'):
# apply gradients