theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-04-23 14:26:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reorder optimizer and logger calls in training loop (#1404)

Browse Source
This commit is contained in:
Qifan Zhang
2026-04-21 13:45:09 +08:00
committed by GitHub
parent 079e51c9f3
commit 5c89a15b9a
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
diffsynth/diffusion/runner.py
View File

@@ -33,15 +33,15 @@ def launch_training_task(
for epoch_id in range(num_epochs): for epoch_id in range(num_epochs):
for data in tqdm(dataloader): for data in tqdm(dataloader):
with accelerator.accumulate(model): with accelerator.accumulate(model):
optimizer.zero_grad()
if dataset.load_from_cache: if dataset.load_from_cache:
loss = model({}, inputs=data) loss = model({}, inputs=data)
else: else:
loss = model(data) loss = model(data)
accelerator.backward(loss) accelerator.backward(loss)
optimizer.step() optimizer.step()
model_logger.on_step_end(accelerator, model, save_steps, loss=loss)
scheduler.step() scheduler.step()
optimizer.zero_grad()
model_logger.on_step_end(accelerator, model, save_steps, loss=loss)
if save_steps is None: if save_steps is None:
model_logger.on_epoch_end(accelerator, model, epoch_id) model_logger.on_epoch_end(accelerator, model, epoch_id)
model_logger.on_training_end(accelerator, model, save_steps) model_logger.on_training_end(accelerator, model, save_steps)