class_docstr: After every model forward, verify the loss is not infinite.
  Called for both training/test phase.
class_docstr_tok: ["After","every","model","forward",",","verify","the","loss","is","not","infinite",".","Called","for","both","training\/test","phase","."]
docstr: Called each time a model forward is done and make sure that
  the model forward output is not NaN. If we encounter NaN as the model
  output, we checkpoint the model to enable debugging and also checkpoint
  the model input sample, model output.
docstr_tok: ["Called","each","time","a","model","forward","is","done","and","make","sure","that","the","model","forward","output","is","not","NaN",".","If","we","encounter","NaN","as","the","model","output",",","we","checkpoint","the","model","to","enable","debugging","and","also","checkpoint","the","model","input","sample",",","model","output","."]
code: def on_forward(self, task: "tasks.ClassyTask") -> None:
      """
      Called each time a model forward is done and make sure that
      the model forward output is not NaN. If we encounter NaN as the model
      output, we checkpoint the model to enable debugging and also checkpoint
      the model input sample, model output.
      """
      # check the model output is not NaN.
      model_output = task.last_batch.model_output
      has_nan = model_output_has_nan(model_output)
  
      if has_nan:
          _, dist_rank = get_machine_local_and_dist_rank()
          logging.info(
              f"Infinite Model output or NaN at iteration={task.iteration}."
          )
  
          # TODO - this code was broken during a refactoring: improve it
          from vissl.hooks.log_hooks import LogLossMetricsCheckpointHook
  
          LogLossMetricsCheckpointHook.checkpoint_model(
              task,
              world_size=self.world_size,
              mode_frequency=1,
              mode_num=task.iteration,
              mode="iteration",
          )
          model_output_file = f"{task.checkpoint_folder}/rank{dist_rank}_model_output.pth"
          input_sample_file = f"{task.checkpoint_folder}/rank{dist_rank}_input_sample.pth"
          with PathManager.open(model_output_file, "wb") as fwrite:
              torch.save(model_output, fwrite)
          with PathManager.open(input_sample_file, "wb") as fwrite:
              torch.save(task.last_batch.sample, fwrite)
          logging.info(f"Saved model output: {model_output_file}")
          logging.info(f"Saved model input: {input_sample_file}")
  
imports: import logging
  import time
  import torch
  from classy_vision import tasks
  from classy_vision.generic.profiler import compute_activations, compute_flops, count_params
  from classy_vision.hooks.classy_hook import ClassyHook
  from fvcore.common.file_io import PathManager
  from vissl.data import AirstoreDataset, GenericSSLDataset
  from vissl.models.model_helpers import model_output_has_nan
  from vissl.utils.env import get_machine_local_and_dist_rank
label_desc: Use image node_id 2 for calling the CheckNanModelOutputHook obj's underlying member method code with example usage: obj.on_forward(task) without return types