Add the --mps_only option. (#1432)

Summary: Pull Request resolved: #1432

Reviewed By: davidberard98

Differential Revision: D43587178

Pulled By: xuzhao9

fbshipit-source-id: d10ed14aca5a17ceea29ac680351120477fbcd41

conftest.py

@@ -18,6 +18,8 @@ def pytest_addoption(parser):
  help="Run benchmarks on cpu only and ignore machine configuration checks")
  parser.addoption("--cuda_only", action='store_true',
  help="Run benchmarks on cuda only and ignore machine configuration checks")
+ parser.addoption("--mps_only", action='store_true',
+ help="Run benchmarks on mps only and ignore machine configuration checks")
 
 def set_fuser(fuser):
  if fuser == "te":

test_bench.py

@@ -34,6 +34,9 @@ def pytest_generate_tests(metafunc):
  if metafunc.config.option.cuda_only:
  devices = ['cuda']
 
+ if metafunc.config.option.mps_only:
+ devices = ['mps']
+
  if metafunc.cls and metafunc.cls.__name__ == "TestBenchNetwork":
  paths = _list_model_paths()
  metafunc.parametrize(