zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-31 11:54:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix gradient accumulation averaging

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-10-10 02:45:26 -07:00
parent 3c587ed618
commit fc88e3b0d0
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
flux/dreambooth.py
View File

@ -274,7 +274,11 @@ if __name__ == "__main__":
loss, grads = nn.value_and_grad(flux.flow, flux.training_loss)(
x, t5_feat, clip_feat, guidance
)
grads = tree_map(lambda a, b: a + b, prev_grads, grads)
grads = tree_map(
lambda a, b: (a + b) / args.grad_accumulate,
prev_grads,
grads,
)
grads = average_gradients(grads)
optimizer.update(flux.flow, grads)