diff --git a/torch_npu/csrc/npu/Module.cpp b/torch_npu/csrc/npu/Module.cpp index c66f2a8d52392bf0c3f3ba05457c77a5bec06f5e..79562439ffd01e18e51e428f43e727c918068749 100644 --- a/torch_npu/csrc/npu/Module.cpp +++ b/torch_npu/csrc/npu/Module.cpp @@ -1651,7 +1651,20 @@ PyObject* THNPModule_aclnn_reselect_static_kernel(PyObject* self, PyObject* noar HANDLE_TH_ERRORS NPUStatus ret = c10_npu::emptyAllNPUStream(); TORCH_CHECK(ret == NPU_STATUS_SUCCESS, "Failed to empty NPU task queue, ret:", ret, PTA_ERROR(ErrCode::INTERNAL)); - c10_npu::opapi::ReselectStaticKernel(); + + static const auto task_queue_enable = c10_npu::option::OptionsManager::GetTaskQueueEnable(); + if (task_queue_enable == 2) { + auto acl_call = []()->int { + c10_npu::opapi::ReselectStaticKernel(); + return 0; + }; + at_npu::native::OpCommand::RunOpApiV2("reselect_static_kernel", acl_call); + NPUStatus ret = c10_npu::emptyAllNPUStream(); + TORCH_CHECK(ret == NPU_STATUS_SUCCESS, "Failed to empty NPU task queue, ret:", ret, PTA_ERROR(ErrCode::INTERNAL)); + } else { + c10_npu::opapi::ReselectStaticKernel(); + } + Py_RETURN_NONE; END_HANDLE_TH_ERRORS }