class_docstr: Hook called after every forward pass (to check training doesn't give NaN),
  after every step and at the end of epoch (to check if the model should be checkpointed)
  and print the meters values at the end of every phase.
class_docstr_tok: ["Hook","called","after","every","forward","pass","(","to","check","training","does","n't","give","NaN",")",",","after","every","step","and","at","the","end","of","epoch","(","to","check","if","the","model","should","be","checkpointed",")","and","print","the","meters","values","at","the","end","of","every","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}."
          )
          self.checkpoint_model(
              task,
              mode_frequency=1,
              mode_num=task.iteration,
              mode="iteration",
              world_size=self.world_size,
          )
          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 g_pathmgr.open(model_output_file, "wb") as fwrite:
              torch.save(model_output, fwrite)
          with g_pathmgr.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 atexit
  import datetime
  import json
  import logging
  import time
  from typing import List, Optional
  import torch
  from classy_vision import meters, tasks
  from classy_vision.generic.distributed_util import get_rank, is_primary
  from classy_vision.hooks.classy_hook import ClassyHook
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from iopath.common.file_io import g_pathmgr
  from vissl.models.model_helpers import model_output_has_nan
  from vissl.utils.checkpoint import CheckpointWriter, is_checkpoint_phase
  from vissl.utils.env import get_machine_local_and_dist_rank
  from vissl.utils.io import save_file
  from vissl.utils.logger import log_gpu_stats
  from vissl.utils.perf_stats import PerfStats
label_desc: Use image node_id 3 for calling the LogLossMetricsCheckpointHook obj's underlying member method code with example usage: obj.on_forward(task) without return types