Optimize AddTakeGrad Tensor Sum (apache#17906)

mk-61 · Apr 7, 2020 · 886b90b · 886b90b
1 parent 44486fc
commit 886b90b
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/3rdparty/mshadow/mshadow/tensor_cpu-inl.h b/3rdparty/mshadow/mshadow/tensor_cpu-inl.h
@@ -503,17 +503,20 @@ template<bool clip, typename IndexType, typename DType>
 inline void AddTakeGrad(Tensor<cpu, 2, DType> dst,
                         const Tensor<cpu, 1, IndexType>& index,
                         const Tensor<cpu, 2, DType> &src) {
-  const int K = dst.shape_[0];
+  const index_t K = dst.shape_[0];
+  const index_t C = dst.shape_[1];
   for (index_t y = 0; y < index.size(0); ++y) {
-    int j = index[y];
+    index_t j = index[y];
     if (clip) {
       if (j <= 0) j = 0;
       else if (j >= K) j = K - 1;
     } else {
       j %= K;
       if (j < 0) j += K;
     }
-    dst[j] += src[y];
+    for (index_t i = 0; i < C; ++i) {
+      dst[j][i] += src[y][i];
+    }
   }
 }