OSSInnovation
/
mindspore
0
0
Fork 26
Code Issues Pull Requests 7 Releases Wiki Activity

!14371 GPU update addn 

From: @VectorSL
Reviewed-by: @cristoval,@limingqi107
Signed-off-by: @limingqi107
This commit is contained in:
mindspore-ci-bot 2021-03-30 18:32:46 +08:00 committed by Gitee
parent 46ed0928b1 8407ccf455
commit 1a07d46a16
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mindspore/ccsrc/backend/kernel_compiler/gpu/math/addn_gpu_kernel.h
View File

@ -108,20 +108,20 @@ class AddNGpuFwdKernel : public GpuKernel {
for (size_t i = input_shape.size(); i < 4; i++) {
(void)input_shape.insert(input_shape.begin(), 1);
}
int dimA[4];
std::vector<int> dimA;
for (size_t i = 0; i < input_shape.size(); i++) {
dimA[i] = SizeToInt(input_shape[i]);
dimA.push_back(SizeToInt(input_shape[i]));
}
auto input_format = AnfAlgo::GetInputFormat(kernel_node, 0);
if (input_format == kOpFormat_NHWC) {
CHECK_CUDNN_RET_WITH_EXCEPT(kernel_node_,
cudnnSetTensorNdDescriptorEx(input_descriptor_, CUDNN_TENSOR_NHWC, cudnn_data_type_,
SizeToInt(input_shape.size()), dimA),
SizeToInt(input_shape.size()), dimA.data()),
"cudnnSetTensorNdDescriptor failed");
} else {
CHECK_CUDNN_RET_WITH_EXCEPT(kernel_node_,
cudnnSetTensorNdDescriptorEx(input_descriptor_, CUDNN_TENSOR_NCHW, cudnn_data_type_,
SizeToInt(input_shape.size()), dimA),
SizeToInt(input_shape.size()), dimA.data()),
"cudnnSetTensorNdDescriptor failed");
}
InitSizeLists();