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: Checkpoint model. Can be called in 3 possible scenarios:
  1. If training becomes NaN, then we checkpoint the model to facilitate debugging
  2. After every N epochs (CHECKPOINT_FREQ), model state is checkpointed.
  3. If user wants to checkpoint during the epoch (ie. after every few training
     iterations, the model state is checkpointed.)
  
  Args:
      task: Self-supervision task that hold information about training iteration,
            epoch number etc.
      train_phase_idx (int): current training phase number. Starts from 0
      mode_frequency (int): mode can be "phase" or "iteration". Frequency
                            of checkpointing for the given mode
      mode_num (int): for the checkpointing mode (phase or iteration), the number
                      of phase or iteration at which checkpointing is being done
docstr_tok: ["Checkpoint","model",".","Can","be","called","in","3","possible","scenarios",":","1",".","If","training","becomes","NaN",",","then","we","checkpoint","the","model","to","facilitate","debugging","2",".","After","every","N","epochs","(","CHECKPOINT_FREQ",")",",","model","state","is","checkpointed",".","3",".","If","user","wants","to","checkpoint","during","the","epoch","(","ie",".","after","every","few","training","iterations",",","the","model","state","is","checkpointed",".",")","Args",":","task",":","Self-supervision","task","that","hold","information","about","training","iteration",",","epoch","number","etc",".","train_phase_idx","(","int",")",":","current","training","phase","number",".","Starts","from","0","mode_frequency","(","int",")",":","mode","can","be","``","phase","''","or","``","iteration","''",".","Frequency","of","checkpointing","for","the","given","mode","mode_num","(","int",")",":","for","the","checkpointing","mode","(","phase","or","iteration",")",",","the","number","of","phase","or","iteration","at","which","checkpointing","is","being","done"]
code: def checkpoint_model(
      task: "tasks.ClassyTask",
      world_size: int,
      mode_frequency: int,
      mode_num: int,
      mode: str = "phase",
  ):
      """
      Checkpoint model. Can be called in 3 possible scenarios:
      1. If training becomes NaN, then we checkpoint the model to facilitate debugging
      2. After every N epochs (CHECKPOINT_FREQ), model state is checkpointed.
      3. If user wants to checkpoint during the epoch (ie. after every few training
         iterations, the model state is checkpointed.)
  
      Args:
          task: Self-supervision task that hold information about training iteration,
                epoch number etc.
          train_phase_idx (int): current training phase number. Starts from 0
          mode_frequency (int): mode can be "phase" or "iteration". Frequency
                                of checkpointing for the given mode
          mode_num (int): for the checkpointing mode (phase or iteration), the number
                          of phase or iteration at which checkpointing is being done
      """
      phase_idx = task.phase_idx
      train_phase_idx = task.train_phase_idx
      # num_train_phases = num_epochs * num_phases_per_epoch
      # For OSS use, num_train_phases will be equal to num_epochs
      num_train_phases = task.num_train_phases
  
      # check if we need to checkpoint this phase
      is_checkpointing_phase = is_checkpoint_phase(
          mode_num,
          mode_frequency,
          train_phase_idx,
          num_train_phases,
          mode,
      )
      is_final_train_phase = (
          (train_phase_idx == (num_train_phases - 1))
          and task.train
          and mode == "phase"
      )
  
      # handle checkpoint:
      if task.train and (
          is_final_train_phase or is_checkpointing_phase
      ):
          #  - if sharded state consolidate the state
          # /!\ All the ranks have to participate
          if (
              hasattr(task.optimizer, "consolidate_state_dict")
              and mode != "phase"
          ):
              logging.info(
                  f"[{mode}: {mode_num}] Consolidating sharded state on all replicas"
              )
              task.optimizer.consolidate_state_dict()
  
          # Depending on whether we are in FSDP mode or not
          # - save the checkpoint on the primary rank
          # - save the sharded checkpoint on all ranks
          if is_primary() or isinstance(task.base_model, FSDP):
              checkpoint_folder = task.checkpoint_folder
              logging.info(
                  f"[{mode}: {mode_num}] Saving checkpoint to {checkpoint_folder}"
              )
              model_state_dict = task.get_classy_state()
  
              # phase_idx is already incremented at the beginning of phase but if we
              # are checkpointing at an iteration in the middle of phase, we should not
              # save the incremented phase_idx as it will incorrectly assume that model
              # trained for that phase already.
              if mode == "iteration":
                  model_state_dict["phase_idx"] = (
                      model_state_dict["phase_idx"] - 1
                  )
                  if task.train:
                      train_phase_idx = train_phase_idx - 1
                      model_state_dict[
                          "train_phase_idx"
                      ] = train_phase_idx
                  restart_phase = phase_idx - 1
                  restart_iteration = task.iteration
  
              # When loading from a phase checkpoint:
              else:
                  restart_phase = phase_idx
                  restart_iteration = task.iteration
  
              if task.ema_model is not None:
                  model_state_dict[
                      "ema_model"
                  ] = task.ema_model.module.get_classy_state()
  
                  model_state_dict["ema_meters"] = [
                      meter.get_classy_state()
                      for meter in task.ema_meters
                  ]
  
              # Content of the loss to be saved
              # - for DDP model, it contains the whole loss
              # - for FSDP model, loss can override state_dict() to
              #   save only the part related to the shard
              loss_state_dict = task.loss.state_dict()
  
              # Content of the checkpoint to be saved:
              # - for DDP model, it contains the whole model
              # - for FSDP model, it contains a shard of the model
              checkpoint_content = {
                  "phase_idx": restart_phase,
                  "iteration": restart_iteration,
                  "loss": loss_state_dict,
                  "iteration_num": task.local_iteration_num,
                  "train_phase_idx": train_phase_idx,
                  "classy_state_dict": model_state_dict,
              }
  
              # Saving the checkpoint:
              # - for DDP model, the primary save the whole model
              # - for FSDP model, each rank saves its own shard
              checkpoint_writer = CheckpointWriter(
                  checkpoint_folder=checkpoint_folder,
                  is_final_train_phase=is_final_train_phase,
                  mode=mode,
                  mode_num=mode_num,
                  backend=task.config["CHECKPOINT"]["BACKEND"],
              )
              if isinstance(task.base_model, FSDP):
                  _, rank = get_machine_local_and_dist_rank()
                  checkpoint_writer.save_sharded_checkpoint(
                      content=checkpoint_content,
                      shard_rank=rank,
                      world_size=world_size,
                  )
              else:
                  checkpoint_writer.save_consolidated_checkpoint(
                      checkpoint_content
                  )
  
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 6 for calling the LogLossMetricsCheckpointHook obj's underlying member method code with example usage: obj.checkpoint_model(task, world_size, mode_frequency, mode_num, mode) without return types