`nemo_rl.utils.grad_norm`#

Module Contents#

Functions#

warn_if_inf_grad_norm

Print a diagnostic warning on rank 0 if grad_norm is infinite.

Data#

_INF_GRAD_NORM_WARNING

API#

nemo_rl.utils.grad_norm._INF_GRAD_NORM_WARNING#: ‘strip(…)’

nemo_rl.utils.grad_norm.warn_if_inf_grad_norm(grad_norm) → None#

Print a diagnostic warning on rank 0 if grad_norm is infinite.

Accepts a float, int, or scalar torch.Tensor. Silently ignores None (e.g. when there are no trainable parameters on a rank).