!23254 Fix multinomial cpu kernel with two-dimension input

Merge pull request !23254 from chengang/fix_cpu_multinomial
This commit is contained in:
i-robot 2021-09-11 02:44:51 +00:00 committed by Gitee
commit e2950f4814
1 changed files with 4 additions and 3 deletions

View File

@ -78,10 +78,11 @@ bool MultinomialGpuKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
} }
// Normalize the cumulative array. // Normalize the cumulative array.
float sum = cumulative_value[num_col - 1]; float sum = cumulative_value[(i + 1) * num_col - 1];
if (sum != 0) { if (sum != 0) {
for (int k = 0; k < num_col; ++k) { for (int k = 0; k < num_col; ++k) {
cumulative_value[k] /= sum; size_t index = i * num_col + k;
cumulative_value[index] /= sum;
} }
} }
@ -113,7 +114,7 @@ bool MultinomialGpuKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
begin = pivot + 1; begin = pivot + 1;
} }
} }
output[i * num_col + n] = begin; output[i * num_sample + n] = begin;
} }
} }
return true; return true;