已合并
[fix]meta_check-on_rebuild-npu_tensor #32899
cuiduo创建于 4月1日
[fix]meta_check-on_rebuild-npu_tensor #32899
已合并
共 2 个文件变更+16-2
| @@ -7,6 +7,7 @@ from torch.testing._internal.common_utils import ( | |||
| 7 | from torch._prims.rng_prims import register_run_and_save_rng_state_op, run_and_save_rng_state | 7 | from torch._prims.rng_prims import register_run_and_save_rng_state_op, run_and_save_rng_state |
| 8 | from testutils import TestUtils | 8 | from testutils import TestUtils |
| 9 | import torch_npu | 9 | import torch_npu |
| 10 | +import torch_npu._inductor | ||
| 10 | 11 | ||
| 11 | 12 | ||
| 12 | class TestRNGPrims(TestUtils): | 13 | class TestRNGPrims(TestUtils): |
| @@ -22,11 +22,24 @@ def _rebuild_npu_tensor(storage, storage_offset, size, stride, requires_grad, ba | |||
| 22 | "please use newer torch to re-store the weight file." | 22 | "please use newer torch to re-store the weight file." |
| 23 | ) | 23 | ) |
| 24 | se._warn_legacy_serialization(warn_massages, "oldfile") | 24 | se._warn_legacy_serialization(warn_massages, "oldfile") |
| 25 | - tensor = torch.tensor([], dtype=storage.dtype, device=storage._untyped_storage.device) | 25 | + tensor = torch.empty( |
| 26 | + (0,), | ||
| 27 | + dtype=storage.dtype, | ||
| 28 | + device=storage._untyped_storage.device, | ||
| 29 | + requires_grad=requires_grad, | ||
| 30 | + ) | ||
| 26 | tensor.set_(storage, storage_offset, size, stride) | 31 | tensor.set_(storage, storage_offset, size, stride) |
| 27 | tensor.requires_grad = requires_grad | 32 | tensor.requires_grad = requires_grad |
| 28 | tensor._backward_hooks = backward_hooks | 33 | tensor._backward_hooks = backward_hooks |
| 29 | - if not se.RE_MAP_CPU: | 34 | + target_device = torch.device("cpu") if se.RE_MAP_CPU else torch.device("npu") |
| 35 | + is_fake_mode = ( | ||
| 36 | + hasattr(torch, "_guards") | ||
| 37 | + and torch._guards.detect_fake_mode(None) is not None | ||
| 38 | + ) | ||
| 39 | + | ||
| 40 | + if is_fake_mode: | ||
| 41 | + tensor.fake_device = target_device | ||
| 42 | + elif not se.RE_MAP_CPU: | ||
| 30 | if isinstance(npu_storage_info, bool): | 43 | if isinstance(npu_storage_info, bool): |
| 31 | tensor = tensor.npu() | 44 | tensor = tensor.npu() |
| 32 | else: | 45 | else: |