已合并
[fix]trace_fix #31328
cuiduo创建于 3月4日
[fix]trace_fix #31328
已合并
共 2 个文件变更+5-0
| @@ -10,6 +10,7 @@ from torch.testing._internal.common_utils import ( | |||
| 10 | ) | 10 | ) |
| 11 | from testutils import TestUtils | 11 | from testutils import TestUtils |
| 12 | import torch_npu | 12 | import torch_npu |
| 13 | +import torch_npu._inductor | ||
| 13 | 14 | ||
| 14 | 15 | ||
| 15 | class TestWrapTriton(TestUtils): | 16 | class TestWrapTriton(TestUtils): |
| @@ -44,6 +44,9 @@ torch_non_c_binding_in_graph_functions_npu = dict.fromkeys( | |||
| 44 | "torch.npu._memory_viz._frame_fmt", | 44 | "torch.npu._memory_viz._frame_fmt", |
| 45 | "torch.npu.amp.autocast_mode.custom_bwd", | 45 | "torch.npu.amp.autocast_mode.custom_bwd", |
| 46 | "torch.npu.amp.autocast_mode.custom_fwd", | 46 | "torch.npu.amp.autocast_mode.custom_fwd", |
| 47 | + "torch.npu._get_current_allocator", | ||
| 48 | + "torch.npu.is_bf16_supported", | ||
| 49 | + | ||
| 47 | ], | 50 | ], |
| 48 | TorchInGraphFunctionVariable, | 51 | TorchInGraphFunctionVariable, |
| 49 | ) | 52 | ) |
| @@ -66,6 +69,7 @@ torch_c_binding_in_graph_functions_npu = dict.fromkeys( | |||
| 66 | "torch_npu._C._npu_setDevice", | 69 | "torch_npu._C._npu_setDevice", |
| 67 | "torch_npu._C._npu_setMemoryFraction", | 70 | "torch_npu._C._npu_setMemoryFraction", |
| 68 | "torch_npu._C._npu_synchronize", | 71 | "torch_npu._C._npu_synchronize", |
| 72 | + "torch_npu._C._npu_resetAccumulatedMemoryStats", | ||
| 69 | ], | 73 | ], |
| 70 | TorchInGraphFunctionVariable, | 74 | TorchInGraphFunctionVariable, |
| 71 | ) | 75 | ) |