format code

pytorch-labs · Nov 27, 2024 · 0a498cc · 0a498cc
1 parent 5cb7c02
commit 0a498cc
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/tritonbench/utils/triton_op.py b/tritonbench/utils/triton_op.py
@@ -1100,14 +1100,17 @@ def _init_extra_metrics() -> Dict[str, Any]:
                     for metric_name, metric_value in nsys_analyzer_results.items():
                         metrics.extra_metrics[metric_name] = metric_value
             if "nsys_gpu_speedup" in self.required_metrics:
+                baseline_nsys_gpu_kernel_sum = (
+                    self.baseline_metrics.extra_metrics.get("nsys_gpu_kernel_sum", None)
+                    if self.baseline_metrics
+                    else None
+                )
+                current_nsys_gpu_kernel_sum = metrics.extra_metrics.get(
+                    "nsys_gpu_kernel_sum", None
+                )
                 metrics.nsys_gpu_speedup = (
-                    self.baseline_metrics.nsys_gpu_kernel_sum
-                    / metrics.nsys_gpu_kernel_sum
-                    if (
-                        self.baseline_metrics
-                        and self.baseline_metrics.nsys_gpu_kernel_sum
-                    )
-                    and metrics.nsys_gpu_kernel_sum
+                    baseline_nsys_gpu_kernel_sum / current_nsys_gpu_kernel_sum
+                    if baseline_nsys_gpu_kernel_sum and current_nsys_gpu_kernel_sum
                     else None
                 )
             if "kineto_trace" in self.required_metrics: