zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-01 12:49:50 +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
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)( loss, grads = nn.value_and_grad(flux.flow, flux.training_loss)(
x, t5_feat, clip_feat, guidance 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) grads = average_gradients(grads)
optimizer.update(flux.flow, grads) optimizer.update(flux.flow, grads)