wehub-resource-sync
35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
# Copyright FunASR (https://github.com/modelscope/FunASR). All Rights Reserved.
|
|
# MIT License (https://opensource.org/licenses/MIT)
|
|
|
|
"""Device helpers for Fun-ASR-Nano runtime paths."""
|
|
|
|
_SUPPORTED_AUTOCAST_DEVICE_TYPES = {"cuda", "xpu", "mps", "npu"}
|
|
|
|
|
|
def _device_type_from_value(device):
|
|
"""Resolve a device type without requiring optional backend registration."""
|
|
if device is None:
|
|
return "cpu"
|
|
|
|
device_type = getattr(device, "type", None)
|
|
if device_type:
|
|
return str(device_type).lower()
|
|
|
|
if isinstance(device, str):
|
|
return device.split(":", 1)[0].lower()
|
|
|
|
return str(device).split(":", 1)[0].lower()
|
|
|
|
|
|
def resolve_autocast_device_type(device):
|
|
"""Return the torch.autocast device_type for a Fun-ASR-Nano device.
|
|
|
|
PyTorch builds without torch_npu may reject ``torch.device("npu:0")`` before
|
|
torch_npu registers the backend. Parse strings directly so NPU requests do
|
|
not fall back to CPU autocast, which only supports bf16 and caused #3034.
|
|
"""
|
|
device_type = _device_type_from_value(device)
|
|
if device_type in _SUPPORTED_AUTOCAST_DEVICE_TYPES:
|
|
return device_type
|
|
return "cpu"
|