Use torch.accelerator in infer_device so non-CUDA accelerators are detected - #3080

Problem

speechbrain.utils.distributed.infer_device() guesses the running device with torch.cuda.is_available(), falling back to "cpu" on any other accelerator. On machines with a non-CUDA accelerator (e.g. Ascend NPU via torch_npu, or any backend registered through the torch.accelerator API), this forces Brain and the pretrained interfaces (speechbrain.inference.interfaces) onto the CPU even though an accelerator is present, and multi-process runs all pile onto the CPU instead of using one accelerator per rank.

Root cause

def infer_device() -> str:
    if torch.cuda.is_available():
        device = "cuda"
        ...
    else:
        device = "cpu"
    return device

Only CUDA is considered; every other accelerator is invisible to the guess.

Fix

Use the device-agnostic PyTorch accelerator API, which returns "cuda" on CUDA machines (identical behavior) and the actual accelerator type (e.g. "npu", "mps") where available:

if torch.accelerator.is_available():
    device = str(torch.accelerator.current_accelerator())
    local_rank = get_local_rank()
    if local_rank is not None:
        device += f":{local_rank}"
else:
    device = "cpu"

current_accelerator() returns a torch.device on newer versions, so the result is stringified before appending the rank suffix (the function's return type stays str).

Adds unit tests for both the plain and LOCAL_RANK paths.

Note: torch.accelerator is available since PyTorch 2.4, while SpeechBrain pins torch>=2.1.0 — happy to guard this with a getattr fallback if you prefer keeping 2.1–2.3 compatibility.

Verification

Verified on Ascend 910B (aarch64, torch 2.14.0a0 + torch_npu 2.14.0):

• old code returns "cpu" while an NPU is available (the bug)
• new code returns "npu"; torch.nn.Linear(2, 2).to("npu") plus a forward pass run correctly
• with LOCAL_RANK=1 the returned device is "npu:1"
• on CUDA machines the returned value is unchanged ("cuda" / "cuda:<rank>")