Upload edit\Qwen3-TTS-test\.venv\Lib\site-packages\torch\mtia\_utils.py with huggingface_hub
Browse files
edit//Qwen3-TTS-test//.venv//Lib//site-packages//torch//mtia//_utils.py
ADDED
|
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from typing import Any
|
| 2 |
+
|
| 3 |
+
import torch
|
| 4 |
+
|
| 5 |
+
# The _get_device_index has been moved to torch.utils._get_device_index
|
| 6 |
+
from torch._utils import _get_device_index as _torch_get_device_index
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def _get_device_index(
|
| 10 |
+
device: Any, optional: bool = False, allow_cpu: bool = False
|
| 11 |
+
) -> int:
|
| 12 |
+
r"""Get the device index from :attr:`device`, which can be a torch.device object, a Python integer, or ``None``.
|
| 13 |
+
|
| 14 |
+
If :attr:`device` is a torch.device object, returns the device index if it
|
| 15 |
+
is a MTIA device. Note that for a MTIA device without a specified index,
|
| 16 |
+
i.e., ``torch.device('mtia')``, this will return the current default MTIA
|
| 17 |
+
device if :attr:`optional` is ``True``. If :attr:`allow_cpu` is ``True``,
|
| 18 |
+
CPU devices will be accepted and ``-1`` will be returned in this case.
|
| 19 |
+
|
| 20 |
+
If :attr:`device` is a Python integer, it is returned as is.
|
| 21 |
+
|
| 22 |
+
If :attr:`device` is ``None``, this will return the current default MTIA
|
| 23 |
+
device if :attr:`optional` is ``True``.
|
| 24 |
+
"""
|
| 25 |
+
if isinstance(device, int):
|
| 26 |
+
return device
|
| 27 |
+
if isinstance(device, str):
|
| 28 |
+
device = torch.device(device)
|
| 29 |
+
if isinstance(device, torch.device):
|
| 30 |
+
if allow_cpu:
|
| 31 |
+
if device.type not in ["mtia", "cpu"]:
|
| 32 |
+
raise ValueError(f"Expected a mtia or cpu device, but got: {device}")
|
| 33 |
+
elif device.type != "mtia":
|
| 34 |
+
raise ValueError(f"Expected a mtia device, but got: {device}")
|
| 35 |
+
if not torch.jit.is_scripting():
|
| 36 |
+
if isinstance(device, torch.mtia.device):
|
| 37 |
+
return device.idx
|
| 38 |
+
return _torch_get_device_index(device, optional, allow_cpu)
|