已合并
[fix]storge_faketensor #33822
cuiduo创建于 4月16日
[fix]storge_faketensor #33822
已合并
共 2 个文件变更+16-3
| @@ -2,7 +2,7 @@ import torch | |||
| 2 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | 2 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests |
| 3 | from testutils import TestUtils | 3 | from testutils import TestUtils |
| 4 | import torch_npu | 4 | import torch_npu |
| 5 | - | 5 | +import torch_npu._inductor |
| 6 | 6 | ||
| 7 | class TestSumAdd(TestUtils): | 7 | class TestSumAdd(TestUtils): |
| 8 | def op_calc(self, input_element, dim, input_element2): | 8 | def op_calc(self, input_element, dim, input_element2): |
| @@ -16,11 +16,24 @@ def _rebuild_npu_tensor(storage, storage_offset, size, stride, requires_grad, ba | |||
| 16 | "please use 2.1 and newer torch to re-store the weight file." | 16 | "please use 2.1 and newer torch to re-store the weight file." |
| 17 | ) | 17 | ) |
| 18 | se._warn_legacy_serialization(warn_massages, "oldfile") | 18 | se._warn_legacy_serialization(warn_massages, "oldfile") |
| 19 | - tensor = torch.tensor([], dtype=storage.dtype, device=storage.device) | 19 | + tensor = torch.empty( |
| 20 | + (0,), | ||
| 21 | + dtype=storage.dtype, | ||
| 22 | + device=storage._untyped_storage.device, | ||
| 23 | + requires_grad=requires_grad, | ||
| 24 | + ) | ||
| 20 | tensor.set_(storage, storage_offset, size, stride) | 25 | tensor.set_(storage, storage_offset, size, stride) |
| 21 | tensor.requires_grad = requires_grad | 26 | tensor.requires_grad = requires_grad |
| 22 | tensor._backward_hooks = backward_hooks | 27 | tensor._backward_hooks = backward_hooks |
| 23 | - if not se.RE_MAP_CPU: | 28 | + target_device = torch.device("cpu") if se.RE_MAP_CPU else torch.device("npu") |
| 29 | + is_fake_mode = ( | ||
| 30 | + hasattr(torch, "_guards") | ||
| 31 | + and torch._guards.detect_fake_mode(None) is not None | ||
| 32 | + ) | ||
| 33 | + | ||
| 34 | + if is_fake_mode: | ||
| 35 | + tensor.fake_device = target_device | ||
| 36 | + elif not se.RE_MAP_CPU: | ||
| 24 | if isinstance(npu_storage_info, bool): | 37 | if isinstance(npu_storage_info, bool): |
| 25 | tensor = tensor.npu() | 38 | tensor = tensor.npu() |
| 26 | else: | 39 | else: |