fix unravelindex cpu

This commit is contained in:
yanzhenxiang2020 2023-01-06 15:34:06 +08:00
parent 4e80bf4c72
commit ca3952c79a
1 changed files with 2 additions and 2 deletions

View File

@ -69,7 +69,7 @@ bool UnravelIndexCpuKernelMod::LaunchKernel(const std::vector<AddressPtr> &input
for (size_t j = 0; j < (inputs[0]->size) / sizeof(T); j++) { for (size_t j = 0; j < (inputs[0]->size) / sizeof(T); j++) {
T Quotient = IndicesData[j]; T Quotient = IndicesData[j];
for (int i = SizeToInt((inputs[1]->size) / sizeof(T) - 1); i >= 0; i--) { for (int i = SizeToInt((inputs[1]->size) / sizeof(T) - 1); i >= 0; i--) {
OutputData[IntToSize(i) + j * ((inputs[1]->size) / sizeof(T))] = Quotient % DimsData[IntToSize(i)]; OutputData[IntToSize(i) * ((inputs[0]->size) / sizeof(T)) + IntToSize(j)] = Quotient % DimsData[IntToSize(i)];
Quotient = (Quotient / DimsData[IntToSize(i)]); Quotient = (Quotient / DimsData[IntToSize(i)]);
} }
} }
@ -78,7 +78,7 @@ bool UnravelIndexCpuKernelMod::LaunchKernel(const std::vector<AddressPtr> &input
for (size_t j = start; j < end; j++) { for (size_t j = start; j < end; j++) {
T Quotient = IndicesData[j]; T Quotient = IndicesData[j];
for (int i = SizeToInt((inputs[1]->size) / sizeof(T) - 1); i >= 0; i--) { for (int i = SizeToInt((inputs[1]->size) / sizeof(T) - 1); i >= 0; i--) {
OutputData[IntToSize(i) + j * ((inputs[1]->size) / sizeof(T))] = Quotient % DimsData[IntToSize(i)]; OutputData[IntToSize(i) * ((inputs[0]->size) / sizeof(T)) + IntToSize(j)] = Quotient % DimsData[IntToSize(i)];
Quotient = (Quotient / DimsData[IntToSize(i)]); Quotient = (Quotient / DimsData[IntToSize(i)]);
} }
} }