class_docstr: A context manager to measure flops of a module. Requires PyTorch 1.13+.
  
  Flop count implementation based on
  https://dev-discuss.pytorch.org/t/the-ideal-pytorch-flop-counter-with-torch-dispatch/505
  
  Examples::
  
      >>> import copy
      >>> import torch
      >>> import torchvision.models as models
      >>> from torchtnt.utils.flops import FlopTensorDispatchMode
  
      >>> module = models.resnet18()
      >>> module_input = torch.randn(1, 3, 224, 224)
      >>> with FlopTensorDispatchMode(module) as ftdm:
      >>>     # count forward flops
      >>>     res = module(module_input).mean()
      >>>     flops_forward = copy.deepcopy(ftdm.flop_counts)
  
      >>>     # reset count before counting backward flops
      >>>     ftdm.reset()
      >>>     res.backward()
      >>>     flops_backward = copy.deepcopy(ftdm.flop_counts)
class_docstr_tok: ["A","context","manager","to","measure","flops","of","a","module",".","Requires","PyTorch","1.13+",".","Flop","count","implementation","based","on","https",":","\/\/dev-discuss.pytorch.org\/t\/the-ideal-pytorch-flop-counter-with-torch-dispatch\/505","Examples",":",":",">",">",">","import","copy",">",">",">","import","torch",">",">",">","import","torchvision.models","as","models",">",">",">","from","torchtnt.utils.flops","import","FlopTensorDispatchMode",">",">",">","module","=","models.resnet18","(",")",">",">",">","module_input","=","torch.randn","(","1",",","3",",","224",",","224",")",">",">",">","with","FlopTensorDispatchMode","(","module",")","as","ftdm",":",">",">",">","#","count","forward","flops",">",">",">","res","=","module","(","module_input",")",".mean","(",")",">",">",">","flops_forward","=","copy.deepcopy","(","ftdm.flop_counts",")",">",">",">","#","reset","count","before","counting","backward","flops",">",">",">","ftdm.reset","(",")",">",">",">","res.backward","(",")",">",">",">","flops_backward","=","copy.deepcopy","(","ftdm.flop_counts",")"]
code: def __exit__(self, exc_type, exc_val, exc_tb):
      for hook_handle in self._all_hooks:
          hook_handle.remove()
      super().__exit__(exc_type, exc_val, exc_tb)
  
imports: import logging
  import operator
  from collections import defaultdict
  from functools import reduce
  from numbers import Number
  from typing import Any, Callable, DefaultDict, Dict, List, Tuple
  import torch
  from torch.utils._python_dispatch import TorchDispatchMode
  from torch.utils._pytree import PyTree, tree_map
label_desc: Use image node_id 2 for calling the FlopTensorDispatchMode obj's underlying member method code with example usage: obj.__exit__(exc_type, exc_val, exc_tb) without return types