Skip to content
4 changes: 3 additions & 1 deletion tests/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@

import torch

from bitsandbytes.cextension import HIP_ENVIRONMENT

test_dims_rng = random.Random(42)


Expand All @@ -21,7 +23,7 @@ def get_available_devices():
# If the environment variable is set, use it directly.
return [os.environ["BNB_TEST_DEVICE"]]

devices = ["cpu"]
devices = [] if HIP_ENVIRONMENT else ["cpu"]

if hasattr(torch, "accelerator"):
# PyTorch 2.6+ - determine accelerator using agnostic API.
Expand Down