Commit 9efb2809 authored by Michael Carilli's avatar Michael Carilli
Browse files

Compatibility between skip_step and FusedAdam.step

parent 111ee132
...@@ -122,7 +122,7 @@ def scale_loss(loss, ...@@ -122,7 +122,7 @@ def scale_loss(loss,
# necessary because amp.scale_loss is already creating a temporary scope. # necessary because amp.scale_loss is already creating a temporary scope.
def patch_step(opt, loss_scaler, loss_id): def patch_step(opt, loss_scaler, loss_id):
opt_step = opt.step opt_step = opt.step
def skip_step(): def skip_step(scale=None):
maybe_print(("Gradient overflow. Skipping step, loss scaler " + maybe_print(("Gradient overflow. Skipping step, loss scaler " +
"{} reducing loss scale to {}").format(loss_id, "{} reducing loss scale to {}").format(loss_id,
loss_scaler.loss_scale())) loss_scaler.loss_scale()))
......
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