keep name same with torchbench

pytorch · Oct 10, 2024 · 4a9920d · 4a9920d
1 parent ec68dea
commit 4a9920d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/torchbenchmark/util/triton_op.py b/torchbenchmark/util/triton_op.py
@@ -850,7 +850,7 @@ def _init_extra_metrics() -> Dict[str, Any]:
                 or "gpu_peak_mem" in self.required_metrics
             ):
                 metrics.cpu_peak_mem, _device_id, metrics.gpu_peak_mem = (
-                    self.get_peak_mem(fn, self.tb_args.metrics_memory_usage_backend)
+                    self.get_peak_mem(fn, self.tb_args.metrics_gpu_backend)
                 )
             if not baseline and "accuracy" in self.required_metrics:
                 metrics.accuracy = (

diff --git a/userbenchmark/triton/run.py b/userbenchmark/triton/run.py
@@ -93,7 +93,7 @@ def get_parser(args=None):
         help="Metrics to collect, split with comma. E.g., --metrics latency,tflops,speedup.",
     )
     parser.add_argument(
-        "--metrics-memory-usage-backend",
+        "--metrics-gpu-backend",
         choices=["default", "torch"],
         default="default",
         help="Specify the backend [default, torch] to collect metrics. In all modes, the latency (execution time) is always collected using `time.time_ns()`. The CPU peak memory usage is collected by `psutil.Process()`. In default mode, the GPU peak memory usage is collected by the `nvml` library. In torch mode, the GPU peak memory usage is collected by `torch.cuda.max_memory_allocated()`.",