VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-17 17:24:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge 3f47806719 into 5e366acda4

Browse Source
This commit is contained in:
Dave Lage
2025-10-10 19:26:56 +01:00
committed by GitHub
parent 5e366acda4 3f47806719
commit b9a9c971ba
3 changed files with 266 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
train_network.py
View File

@@ -767,6 +767,9 @@ class NetworkTrainer:
optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)
optimizer_train_fn, optimizer_eval_fn = train_util.get_optimizer_train_eval_fn(optimizer, args)
if hasattr(network, "register_optimizer"):
network.register_optimizer(optimizer)
# prepare dataloader
# strategies are set here because they cannot be referenced in another process. Copy them with the dataset
# some strategies can be None
@@ -1453,6 +1456,8 @@ class NetworkTrainer:
network.update_grad_norms()
if hasattr(network, "update_norms"):
network.update_norms()
if hasattr(network, "update_gradient_ema"):
network.update_gradient_ema()
optimizer.step()
lr_scheduler.step()