未验证 提交 df601667 编写于 作者: F furnace 提交者: GitHub

[NPU] Fix fill_constant and set_value for CANN 5.0.4 (#39635)

* [NPU] Fix fill_constant int64 for CANN 5.0.4 (FillD does not support int64)

* [NPU] fix set_value for CANN 5.0.4
上级 329b095e
......@@ -65,7 +65,7 @@ class FillConstantNPUKernel : public framework::OpKernel<T> {
tensor_value.mutable_data<T>({1}, ctx.GetPlace());
FillNpuTensorWithConstant<T>(&tensor_value, value);
NpuOpRunner runner;
#if (CANN_VERSION_CODE >= 503003)
#if (CANN_VERSION_CODE >= 503003 && CANN_VERSION_CODE < 504001)
runner.SetType("FillD")
.AddInput(tensor_value)
.AddOutput(*out_var)
......
......@@ -174,6 +174,9 @@ class SetValueNPUKernel : public framework::OpKernel<T> {
.AddInput(std::move(index_indices))
.AddInput(val_temp)
.AddOutput(out_temp)
#if (CANN_VERSION_CODE >= 504001)
.AddAttrs({{"use_locking", false}})
#endif
.Run(stream);
}
};
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册