Add MUSA graph-safe op overrides and torch.load map_location translation

README.md

@@ -392,7 +392,7 @@ See `src/torchada/_mappings/` for 400+ mapping rules grouped by API domain.
 
 ```
 # pyproject.toml or requirements.txt
-torchada>=0.1.83
+torchada>=0.1.84
 ```
 
 ### Step 2: Conditional Import

README_CN.md

@@ -376,7 +376,7 @@ if torchada.is_gpu_device(device):  # 在 CUDA 和 MUSA 上都能工作
 
 ```
 # pyproject.toml 或 requirements.txt
-torchada>=0.1.83
+torchada>=0.1.84
 ```
 
 ### 步骤 2：条件导入

benchmarks/benchmark_history.json

@@ -3,7 +3,7 @@
   "description": "Historical benchmark results for torchada performance tracking",
   "results": [
     {
-      "version": "0.1.83",
+      "version": "0.1.84",
       "date": "2026-01-29",
       "platform": "MUSA",
       "pytorch_version": "2.7.1",

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchada"
-version = "0.1.83"
+version = "0.1.84"
 description = "Adapter package for torch_musa to act exactly like PyTorch CUDA"
 readme = "README.md"
 license = {text = "MIT"}

src/torchada/__init__.py

@@ -24,7 +24,7 @@
     from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CUDA_HOME
 """
 
-__version__ = "0.1.83"
+__version__ = "0.1.84"
 
 from . import cuda, utils
 

src/torchada/_patch.py

@@ -180,6 +180,34 @@ def _translate_device(device: Any) -> Any:
     return device
 
 
+def _translate_map_location(map_location: Any) -> Any:
+    if isinstance(map_location, dict):
+        return {
+            source: _translate_device(target)
+            for source, target in map_location.items()
+        }
+    return _translate_device(map_location)
+
+
+def _wrap_torch_load(original_load: Callable) -> Callable:
+    @functools.wraps(original_load)
+    def wrapped_load(*args, **kwargs):
+        if "map_location" in kwargs:
+            kwargs["map_location"] = _translate_map_location(kwargs["map_location"])
+        elif len(args) >= 2:
+            args = args[:1] + (_translate_map_location(args[1]),) + args[2:]
+        return original_load(*args, **kwargs)
+
+    return wrapped_load
+
+
+@patch_function
+def _patch_torch_load():
+    if hasattr(torch.load, "__wrapped__"):
+        return
+    torch.load = _wrap_torch_load(torch.load)
+
+
 def _wrap_to_method(original_to: Callable) -> Callable:
     """Wrap tensor.to() to translate device strings."""