docstr: Make sure FLOPs calculation works for a module consists of linear layers.
code: def test_torch_linear_layer(self) -> None:
      """Make sure FLOPs calculation works for a module consists of linear layers."""
      lnn = torch.nn.Sequential(
          torch.nn.Sequential(
              torch.nn.Linear(10, 70), torch.nn.Linear(70, 5)
          ),
          torch.nn.Linear(5, 1),
      )
      inp = torch.randn(1, 10)
  
      with FlopTensorDispatchMode(lnn) as ftdm:
          self.assertEqual(len(ftdm._all_hooks), 8)
  
          res = lnn(inp)
          self.assertEqual(
              ftdm.flop_counts[""].get("addmm.default", 0)
              + ftdm.flop_counts[""].get("addmm", 0),
              1055,
          )
          self.assertEqual(
              ftdm.flop_counts["0"].get("addmm.default", 0)
              + ftdm.flop_counts["0"].get("addmm", 0),
              1050,
          )
          self.assertEqual(
              ftdm.flop_counts["0.0"].get("addmm.default", 0)
              + ftdm.flop_counts["0.0"].get("addmm", 0),
              700,
          )
          self.assertEqual(
              ftdm.flop_counts["0.1"].get("addmm.default", 0)
              + ftdm.flop_counts["0.1"].get("addmm", 0),
              350,
          )
          self.assertEqual(
              ftdm.flop_counts["1"].get("addmm.default", 0)
              + ftdm.flop_counts["1"].get("addmm", 0),
              5,
          )
          ftdm.reset()
          res.backward()
          self.assertEqual(
              ftdm.flop_counts[""].get("mm.default", 0)
              + ftdm.flop_counts[""].get("mm", 0),
              1410,
          )
          self.assertEqual(
              ftdm.flop_counts["0"].get("mm.default", 0)
              + ftdm.flop_counts["0"].get("mm", 0),
              1400,
          )
          self.assertEqual(
              ftdm.flop_counts["0.0"].get("mm.default", 0)
              + ftdm.flop_counts["0.0"].get("mm", 0),
              700,
          )
          self.assertEqual(
              ftdm.flop_counts["0.1"].get("mm.default", 0)
              + ftdm.flop_counts["0.1"].get("mm", 0),
              700,
          )
          self.assertEqual(
              ftdm.flop_counts["1"].get("mm.default", 0)
              + ftdm.flop_counts["1"].get("mm", 0),
              10,
          )
  
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 2 for calling the ModuleSummaryTest obj's underlying member method code with example usage: obj.test_torch_linear_layer() without return types