Skip to content

GitLab

Commit 3a701586 authored by Mohammad's avatar Mohammad
Browse files

using None for resetting gradients

parent 8ceed7c7
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
megatron/training.py
View file @ 3a701586
...@@ -227,7 +227,7 @@ def backward_step(optimizer, model, loss): ...@@ -227,7 +227,7 @@ def backward_step(optimizer, model, loss):
timers = get_timers() timers = get_timers()
# Backward pass. # Backward pass.
optimizer.zero_grad() optimizer.zero_grad(set_grads_to_None=True)
if args.fp16: if args.fp16:
optimizer.backward(loss, update_master_grads=False) optimizer.backward(loss, update_master_grads=False)
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment