split out flop counting its own method (#125061)

Summary: Modularizing code for reuse by splitting __torch_dispatch__ to move flop counting to its own method.

Test Plan: unit tests

Differential Revision: D56644523

Pull Request resolved: #125061
Approved by: https://github.com/842974287

torch/utils/flop_counter.py

@@ -539,7 +539,9 @@ def __exit__(self, *args):
     def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         kwargs = kwargs if kwargs else {}
         out = func(*args, **kwargs)
-        func_packet = func._overloadpacket
+        return self._count_flops(func._overloadpacket, out, args, kwargs)
+
+    def _count_flops(self, func_packet, out, args, kwargs):
         if func_packet in self.flop_registry:
             flop_count_func = self.flop_registry[func_packet]
             flop_count = flop_count_func(*args, **kwargs, out_val=out)  # type: ignore[operator]