How to inspect whethet there is NaN or Inf in gradients after amp?

autograd
1

Hi, guys,
I want to know how to inspect whethet there is NaN or Inf in gradients after amp?

optimizer.zero_grad()
with autocast():
    output = model(input)
    loss = loss_fn(output, target)
scaler.scale(loss).backward()
# To Insepct Whethet There is Infs or NaNs in Gradients After AMP?
if contain_inf_nan():
    # Do not skip if there is infs or NaNs
    raise RuntimeError
scaler.step(optimizer)
scaler.update()

Your answer or guide will be appreciated!

1 Like