class_docstr: A callback that performs periodic emptying of cuda cache using `torch.cuda.empty_cache <https://pytorch.org/docs/stable/generated/torch.cuda.empty_cache.html#torch.cuda.empty_cache>_.
  
  On different ranks, reserved memory and fragmentation might diverge after several iterations.
  If different ranks trigger de-fragmentation (i.e. cudaFree and redo cudaMalloc later)
  at different times, there will be different stragglers in different iterations, which will
  hurt the performance and will get worse with larger clusters. To avoid this, this callback
  calls empty_cache() at the same cadence across all ranks.
  
  Args:
      step_interval: number of steps to run before emptying cuda cache
class_docstr_tok: ["A","callback","that","performs","periodic","emptying","of","cuda","cache","using","`","torch.cuda.empty_cache","<","https",":","\/\/pytorch.org\/docs\/stable\/generated\/torch.cuda.empty_cache.html","#","torch.cuda.empty_cache",">","_",".","On","different","ranks",",","reserved","memory","and","fragmentation","might","diverge","after","several","iterations",".","If","different","ranks","trigger","de-fragmentation","(","i.e",".","cudaFree","and","redo","cudaMalloc","later",")","at","different","times",",","there","will","be","different","stragglers","in","different","iterations",",","which","will","hurt","the","performance","and","will","get","worse","with","larger","clusters",".","To","avoid","this",",","this","callback","calls","empty_cache","(",")","at","the","same","cadence","across","all","ranks",".","Args",":","step_interval",":","number","of","steps","to","run","before","emptying","cuda","cache"]
code: def on_train_step_end(self, state: State, unit: TTrainUnit) -> None:
      total_num_steps_completed = (
          unit.train_progress.num_steps_completed
      )
      if state.entry_point == EntryPoint.FIT:
          # if fitting, unit should also subclass EvalUnit
          unit_as_eval_unit = cast(TEvalUnit, unit)
          # if fitting, include the num eval steps completed in the total steps completed
          total_num_steps_completed += (
              unit_as_eval_unit.eval_progress.num_steps_completed
          )
  
      if total_num_steps_completed % self._step_interval == 0:
          torch.cuda.empty_cache()
  
file_path: tnt/torchtnt/framework/callbacks/empty_cuda_cache.py
imports: from typing import cast
  import torch
  from torchtnt.framework.callback import Callback
  from torchtnt.framework.state import EntryPoint, State
  from torchtnt.framework.unit import TEvalUnit, TPredictUnit, TTrainUnit
label_desc: Use image node_id 2 for calling the EmptyCudaCache obj's underlying member method code with example usage: obj.on_train_step_end(state, unit) without return types