class_docstr: A helper class that tracks cpu and gpu memory.
  
  This class will silently skip unless `psutil` is available. Install with `pip install psutil`.
  
  When a stage completes, it can pass metrics dict to update with the memory metrics gathered during this stage.
  
  Example :
  
  ```python
  self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics)
  self._memory_tracker.start()
  # code ...
  metrics = {"train_runtime": 10.5}
  self._memory_tracker.stop_and_update_metrics(metrics)
  ```
  
  At the moment GPU tracking is only for `pytorch`, but can be extended to support `tensorflow`.
  
  To understand this class' intricacies please read the documentation of [`~Trainer.log_metrics`].
class_docstr_tok: ["A","helper","class","that","tracks","cpu","and","gpu","memory",".","This","class","will","silently","skip","unless","`","psutil","`","is","available",".","Install","with","`","pip","install","psutil","`",".","When","a","stage","completes",",","it","can","pass","metrics","dict","to","update","with","the","memory","metrics","gathered","during","this","stage",".","Example",":","``","`","python","self._memory_tracker","=","TrainerMemoryTracker","(","self.args.skip_memory_metrics",")","self._memory_tracker.start","(",")","#","code","...","metrics","=","{","``","train_runtime","''",":","10.5","}","self._memory_tracker.stop_and_update_metrics","(","metrics",")","``","`","At","the","moment","GPU","tracking","is","only","for","`","pytorch","`",",","but","can","be","extended","to","support","`","tensorflow","`",".","To","understand","this","class","'","intricacies","please","read","the","documentation","of","[","`","~Trainer.log_metrics","`","]","."]
docstr: stop tracking for the passed stage
docstr_tok: ["stop","tracking","for","the","passed","stage"]
code: def stop(self, stage):
      """stop tracking for the passed stage"""
  
      # deal with nested calls of eval during train - simply ignore those
      if self.cur_stage is not None and self.cur_stage != stage:
          return
  
      # this sends a signal to peak_monitor_func to complete its loop
      self.peak_monitoring = False
  
      # first ensure all objects get collected and their memory is freed
      gc.collect()
  
      if self.torch is not None:
          if torch.cuda.is_available():
              self.torch.cuda.empty_cache()
          elif is_torch_xpu_available():
              self.torch.xpu.empty_cache()
          elif is_torch_npu_available():
              self.torch.npu.empty_cache()
  
      # concepts:
      # - alloc_delta:  the difference of allocated memory between the end and the start
      # - peaked_delta: the difference between the peak memory and the current memory
      # in order to know how much memory the measured code consumed one needs to sum these two
  
      # gpu
      if self.torch is not None:
          if torch.cuda.is_available():
              self.gpu_mem_used_now = self.torch.cuda.memory_allocated()
              self.gpu_mem_used_peak = (
                  self.torch.cuda.max_memory_allocated()
              )
          elif is_torch_xpu_available():
              self.gpu_mem_used_now = self.torch.xpu.memory_allocated()
              self.gpu_mem_used_peak = (
                  self.torch.xpu.max_memory_allocated()
              )
          elif is_torch_npu_available():
              self.gpu_mem_used_now = self.torch.npu.memory_allocated()
              self.gpu_mem_used_peak = (
                  self.torch.npu.max_memory_allocated()
              )
          else:
              raise ValueError("No available GPU device found!")
  
          self.gpu[self.cur_stage] = {
              "begin": self.gpu_mem_used_at_start,
              "end": self.gpu_mem_used_now,
              "alloc": (
                  self.gpu_mem_used_now - self.gpu_mem_used_at_start
              ),
              "peaked": max(
                  0, self.gpu_mem_used_peak - self.gpu_mem_used_now
              ),
          }
  
      # cpu
      self.cpu_mem_used_now = self.cpu_mem_used()
      self.cpu[self.cur_stage] = {
          "begin": self.cpu_mem_used_at_start,
          "end": self.cpu_mem_used_now,
          "alloc": (self.cpu_mem_used_now - self.cpu_mem_used_at_start),
          "peaked": max(
              0, self.cpu_mem_used_peak - self.cpu_mem_used_now
          ),
      }
  
      # reset - cycle finished
      self.cur_stage = None
  
imports: import copy
  import functools
  import gc
  import inspect
  import os
  import random
  import re
  import threading
  import time
  from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union
  import numpy
  from .utils import ExplicitEnum, is_psutil_available, is_tf_available, is_torch_available, is_torch_cuda_available, is_torch_mps_available, is_torch_npu_available, is_torch_tpu_available, is_torch_xpu_available, requires_backends
label_desc: Use image node_id 6 for calling the TrainerMemoryTracker obj's underlying member method code with example usage: obj.stop(stage) without return types