From 9ada83c7c13b90e3c1cab3d0ed1d691d728c94e8 Mon Sep 17 00:00:00 2001
From: Yuanhao Ji <jiyuanhao@apache.org>
Date: Tue, 19 Nov 2024 10:14:39 +0800
Subject: [PATCH] add --device_only arg

---
 conftest.py   | 6 ++++++
 test_bench.py | 3 +++
 2 files changed, 9 insertions(+)

diff --git a/conftest.py b/conftest.py
index 7f61295157..6cbb8a09e2 100644
--- a/conftest.py
+++ b/conftest.py
@@ -38,6 +38,12 @@ def pytest_addoption(parser):
         action="store_true",
         help="Run benchmarks on mps only and ignore machine configuration checks",
     )
+    parser.addoption(
+        "--device_only",
+        action="store",
+        default=None,
+        help="Run benchmarks on the specific device only and ignore machine configuration checks",
+    )
 
 
 def set_fuser(fuser):
diff --git a/test_bench.py b/test_bench.py
index 724b05d201..72db81eb57 100644
--- a/test_bench.py
+++ b/test_bench.py
@@ -31,6 +31,9 @@ def pytest_generate_tests(metafunc):
     if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
         devices.append("mps")
 
+    if device_only := metafunc.config.option.device_only:
+        devices = [device_only]
+
     if metafunc.config.option.cpu_only:
         devices = ["cpu"]