fixed fn_name parameter for BenchmarkOperator.tflops

Summary:
There is a mismatch between the tflops argument fn and the corresponding
parameter. We pass a Callable fn to tflops, which expects a str fn_name.

Reviewed By: sijiac

Differential Revision: D55736424

fbshipit-source-id: 870e515f07e988fd1eadb7c7165e075c56fa21b9

torchbenchmark/util/triton_op.py

@@ -394,7 +394,7 @@ def _do_bench(self,
 
 
     @register_metric()
-    def tflops(self, fn_name: str, example_inputs: Any, metrics: BenchmarkOperatorMetrics) -> List[float]:
+    def tflops(self, fn: Callable, example_inputs: Any, metrics: BenchmarkOperatorMetrics) -> List[float]:
         def _get_flops(self, func: Callable) -> float:
             """By default, use the torch.__dispatch__ based flops counter."""
             from torch.utils.flop_counter import FlopCounterMode
@@ -410,7 +410,6 @@ def work_func():
                 work_func()
             total_flops = sum([v for _, v in flop_counter.flop_counts["Global"].items()])
             return total_flops
-        fn = self._get_bm_func(fn_name)
         if not fn in self._op_flops:
             self._op_flops[fn] = _get_flops(self, fn)
         op_flops = self._op_flops[fn]