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: from abc import ABC, abstractmethod
  from typing import List, Union
  import hydra
  import torch
  from omegaconf import DictConfig, OmegaConf
  from pytorch_lightning import Trainer
  from nemo.collections.asr.metrics.audio import AudioMetricWrapper
  from nemo.core.classes import ModelPT
  from nemo.utils import logging, model_utils
label_desc: Use image node_id 16 for calling the AudioToAudioModel obj's underlying member method code with example usage: obj.on_after_backward() without return types