docstr: zero-out the gradients which any of them is NAN or INF
docstr_tok: ["zero-out","the","gradients","which","any","of","them","is","NAN","or","INF"]
code: def on_after_backward(self):
      """
      zero-out the gradients which any of them is NAN or INF
      """
      super().on_after_backward()
  
      if hasattr(self, "_skip_nan_grad") and self._skip_nan_grad:
          device = next(self.parameters()).device
          valid_gradients = torch.tensor(
              [1], device=device, dtype=torch.float32
          )
  
          # valid_gradients = True
          for param_name, param in self.named_parameters():
              if param.grad is not None:
                  is_not_nan_or_inf = not (
                      torch.isnan(param.grad).any()
                      or torch.isinf(param.grad).any()
                  )
                  if not is_not_nan_or_inf:
                      valid_gradients = valid_gradients * 0
                      break
  
          if torch.distributed.is_initialized():
              torch.distributed.all_reduce(
                  valid_gradients, op=torch.distributed.ReduceOp.MIN
              )
  
          if valid_gradients < 1:
              logging.warning(
                  f"detected inf or nan values in gradients! Setting gradients to zero."
              )
              self.zero_grad()
  
imports: import logging
  from abc import ABC, abstractmethod
  from typing import List
  import torch
  from nemo.core.classes import ModelPT
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.core.classes.exportable import Exportable
  from nemo.core.classes.mixins import AccessMixin
  from nemo.core.utils.neural_type_utils import get_io_names
  from nemo.utils import logging, model_utils
  from nemo.utils.cast_utils import cast_all
label_desc: Use image node_id 7 for calling the ASRModel obj's underlying member method code with example usage: obj.on_after_backward() without return types