Add BNB_TEST_DEVICE env var to manually select device for unit tests

Author: matthewdouglas

tests/helpers.py

@@ -1,6 +1,7 @@
 import functools
 from io import BytesIO
 from itertools import product
+import os
 import random
 from typing import Any
 
@@ -16,6 +17,10 @@
 
 @functools.cache
 def get_available_devices():
+    if "BNB_TEST_DEVICE" in os.environ:
+        # If the environment variable is set, use it directly.
+        return [os.environ["BNB_TEST_DEVICE"]]
+
     devices = ["cpu"]
 
     if hasattr(torch, "accelerator"):