Bugfix

pytorch-labs · Dec 7, 2024 · 8896feb · 8896feb
1 parent e437eb6
commit 8896feb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tritonbench/utils/triton_op.py b/tritonbench/utils/triton_op.py
@@ -1124,8 +1124,8 @@ def _init_extra_metrics() -> Dict[str, Any]:
             if "proton" in self.required_metrics:
                 from tritonbench.components.proton import proton_trace
                 scope_name = fn_name
-                flops = self.flops() if self.has_metric("flops") else None
-                num_bytes = self.bytes() if self.has_metric("bytes") else None
+                flops = self.flops(fn_name, self.example_inputs, metrics) if self.has_metric("flops") else None
+                num_bytes = self.bytes(fn_name, self.example_inputs, metrics) if self.has_metric("bytes") else None
                 proton_trace(self._proton_session_id, scope_name, fn, warmup=warmup, flops=flops, bytes=num_bytes)
             if "best_config" in self.required_metrics:
                 metrics.best_config = self.best_config(fn)