* Fix einsum Bug

apache · May 27, 2020 · 5f83e7d · 5f83e7d
1 parent 3efacd2
commit 5f83e7d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/operator/numpy/np_einsum_op-inl.h b/src/operator/numpy/np_einsum_op-inl.h
@@ -724,7 +724,10 @@ inline void NumpyEinsumProcess(const std::vector<TBlob>& inputs,
         int j = 0;
         for (idim = 0; idim < ndim_iter; ++idim) {
           if (op_axes_arrays[i][idim] == -1 ||
-              opshape[i][op_axes_arrays[i][idim]] == 1) {
+              (iop!=nop && 
+               opshape[i][op_axes_arrays[i][idim]] == 1 && 
+               opshape[iop][op_axes_arrays[iop][idim]] != 1)
+               ){
             remainstride[iop][j++] = iterstride[iop][idim];
           } else {
             opstride[iop][op_axes_arrays[i][idim]] = iterstride[iop][idim];