Add --device_only arg to test_bench.py

[shink]

I noticed that test_bench.py only supports there devices: cpu, cuda and mps. And I think there should be a power to allow users to run benchmarks on other devices (e.g. xpu, npu, etc.). So I added an --device_only arg which can receive a string. Running pytest test_bench.py --device_only foo will make benchmarks to run on the foo device.

BTW, test.py use ACCELERATOR environment variable to custom devices, should we make consistent?

benchmark/test.py

Lines 174 to 175 in 75c0f31

	if device := os.getenv("ACCELERATOR"):
	devices.append(device)

[shink]

Look, device_only arg has lower priority, pytest test_bench.py --cpu_only --device_only cuda will make benchmarks to run on cpu deivices.