docstr: Make sure FLOPs calculation works for a single operations.
code: def test_torch_operations(self) -> None:
      """Make sure FLOPs calculation works for a single operations."""
  
      class TestModule(torch.nn.Module):
          def __init__(self) -> None:
              super().__init__()
  
          def forward(self, x: torch.Tensor) -> torch.Tensor:
              bmm_mat = torch.randn(10, 5, 7)
              mm_mat = torch.randn(7, 3)
              return x.bmm(bmm_mat).matmul(mm_mat)
  
      test_module = TestModule()
      with FlopTensorDispatchMode(test_module) as ftdm:
          inp = torch.randn(10, 4, 5)
          res = test_module(inp)
  
          self.assertEqual(res.shape[0], 10)
          self.assertEqual(res.shape[1], 4)
          self.assertEqual(res.shape[2], 3)
  
          self.assertEqual(
              ftdm.flop_counts[""].get("bmm.default", 0)
              + ftdm.flop_counts[""].get("bmm", 0),
              1400,
          )
          self.assertEqual(
              ftdm.flop_counts[""].get("mm.default", 0)
              + ftdm.flop_counts[""].get("mm", 0),
              840,
          )
  
          inp = torch.randn(10, 4, 5)
          inp = torch.autograd.Variable(inp, requires_grad=True)
  
          ftdm.reset()
          res = test_module(inp)
          res.mean().backward()
  
          self.assertEqual(res.shape[0], 10)
          self.assertEqual(res.shape[1], 4)
          self.assertEqual(res.shape[2], 3)
  
          self.assertEqual(
              ftdm.flop_counts[""].get("bmm.default", 0)
              + ftdm.flop_counts[""].get("bmm", 0),
              2800,
          )
          self.assertEqual(
              ftdm.flop_counts[""].get("mm.default", 0)
              + ftdm.flop_counts[""].get("mm", 0),
              1680,
          )
  
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 1 for calling the ModuleSummaryTest obj's underlying member method code with example usage: obj.test_torch_operations() and returns: x