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

Fix bug with zeroing out gradients when accumulating

Browse Source
This commit is contained in:
Jaret Burkett
2024-09-03 08:29:15 -06:00
parent 121a760c19
commit 5c8fcc8a4e
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -1073,7 +1073,6 @@ class SDTrainer(BaseSDTrainProcess):
# set the weights # set the weights
network.multiplier = network_weight_list network.multiplier = network_weight_list
self.optimizer.zero_grad(set_to_none=True)
# activate network if it exits # activate network if it exits
@@ -1554,8 +1553,8 @@ class SDTrainer(BaseSDTrainProcess):
else: else:
batch_list = [batch] batch_list = [batch]
total_loss = None total_loss = None
self.optimizer.zero_grad()
for batch in batch_list: for batch in batch_list:
self.optimizer.zero_grad(set_to_none=True)
loss = self.train_single_accumulation(batch) loss = self.train_single_accumulation(batch)
if total_loss is None: if total_loss is None:
total_loss = loss total_loss = loss