docstr: Make sure FLOPs calculation works for a resnet18.
code: def test_torch_pretrained_module(self) -> None:
      """Make sure FLOPs calculation works for a resnet18."""
      # pyre-fixme[16]: Module `models` has no attribute `resnet18`.
      mod = models.resnet18()
      inp = torch.randn(1, 3, 224, 224)
      with FlopTensorDispatchMode(mod) as ftdm:
          # Hooks should be 2 * number of modules minus 2 (2 for the model itself)
          self.assertEqual(
              len(ftdm._all_hooks), 2 * len(list(mod.modules())) - 2
          )
          res = mod(inp)
  
          self.assertEqual(
              ftdm.flop_counts[""].get("convolution.default", 0)
              + ftdm.flop_counts[""].get("convolution", 0),
              1813561344,
          )
          self.assertEqual(
              ftdm.flop_counts[""].get("addmm.default", 0)
              + ftdm.flop_counts[""].get("addmm", 0),
              512000,
          )
          self.assertEqual(
              ftdm.flop_counts["conv1"].get("convolution.default", 0)
              + ftdm.flop_counts["conv1"].get("convolution", 0),
              118013952,
          )
          self.assertEqual(
              ftdm.flop_counts["fc"].get("addmm.default", 0)
              + ftdm.flop_counts["fc"].get("addmm", 0),
              512000,
          )
  
          ftdm.reset()
          res.mean().backward()
  
          self.assertEqual(
              ftdm.flop_counts[""].get(
                  "convolution_backward.default", 0
              )
              + ftdm.flop_counts[""].get("convolution_backward", 0),
              3509108736,
          )
          self.assertEqual(
              ftdm.flop_counts[""].get("mm.default", 0)
              + ftdm.flop_counts[""].get("mm", 0),
              1024000,
          )
          self.assertEqual(
              ftdm.flop_counts["layer1"].get(
                  "convolution_backward.default", 0
              )
              + ftdm.flop_counts["layer1"].get(
                  "convolution_backward", 0
              ),
              924844032,
          )
          self.assertEqual(
              ftdm.flop_counts["fc"].get("mm.default", 0)
              + ftdm.flop_counts["fc"].get("mm", 0),
              1024000,
          )
  
file_path: tnt/tests/utils/test_flops.py
imports: import unittest
  import torch
  import torchvision.models
  from torchtnt.utils.flops import FlopTensorDispatchMode
label_desc: Use image node_id 3 for calling the ModuleSummaryTest obj's underlying member method code with example usage: obj.test_torch_pretrained_module() without return types