(no_cpu=False)
| 19 | |
| 20 | @functools.cache |
| 21 | def get_available_devices(no_cpu=False): |
| 22 | if "BNB_TEST_DEVICE" in os.environ: |
| 23 | # If the environment variable is set, use it directly. |
| 24 | device = os.environ["BNB_TEST_DEVICE"] |
| 25 | return [] if no_cpu and device == "cpu" else [device] |
| 26 | |
| 27 | devices = [] if HIP_ENVIRONMENT else ["cpu"] if not no_cpu else [] |
| 28 | |
| 29 | if hasattr(torch, "accelerator"): |
| 30 | # PyTorch 2.6+ - determine accelerator using agnostic API. |
| 31 | if torch.accelerator.is_available(): |
| 32 | devices += [str(torch.accelerator.current_accelerator())] |
| 33 | else: |
| 34 | if torch.cuda.is_available(): |
| 35 | devices += ["cuda"] |
| 36 | |
| 37 | if torch.backends.mps.is_available(): |
| 38 | devices += ["mps"] |
| 39 | |
| 40 | if hasattr(torch, "xpu") and torch.xpu.is_available(): |
| 41 | devices += ["xpu"] |
| 42 | |
| 43 | custom_backend_name = torch._C._get_privateuse1_backend_name() |
| 44 | custom_backend_module = getattr(torch, custom_backend_name, None) |
| 45 | custom_backend_is_available_fn = getattr(custom_backend_module, "is_available", None) |
| 46 | |
| 47 | if custom_backend_is_available_fn and custom_backend_module.is_available(): |
| 48 | devices += [custom_backend_name] |
| 49 | |
| 50 | return devices |
| 51 | |
| 52 | |
| 53 | def torch_save_to_buffer(obj): |
nothing calls this directly
no outgoing calls
no test coverage detected