File tree Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Original file line number Diff line number Diff line change @@ -67,14 +67,16 @@ void TileCompute<T, PType>::Run() {
67
67
tmp_dst_tensor.Resize (out_dims);
68
68
auto tmp_src = tmp_src_tensor.mutable_data <T>();
69
69
auto tmp_dst = tmp_dst_tensor.mutable_data <T>();
70
- for (int i = 0 ; i < in-> dims () .production (); i++) {
70
+ for (int i = 0 ; i < in_dims .production (); i++) {
71
71
tmp_src[i] = in_data[i];
72
72
}
73
+
73
74
for (int i = bcast_dims.size () - 1 ; i >= 0 ; i--) {
74
75
if (bcast_dims[i] > 1 ) {
75
- for (int m = 0 ; m < in_stride[i]; m++) {
76
+ int num = in_stride[1 ] / in_stride[i + 1 ];
77
+ int dst_stride = in_stride[i + 1 ] * bcast_dims[i + 1 ];
78
+ for (int m = 0 ; m < num; m++) {
76
79
for (int j = 0 ; j < bcast_dims[i]; j++) {
77
- int dst_stride = in_stride[i + 1 ] * bcast_dims[i + 1 ];
78
80
std::memcpy (tmp_dst + j * dst_stride + m * bcast_dims[i] * dst_stride,
79
81
tmp_src + m * in_stride[i + 1 ],
80
82
dst_stride * sizeof (T));
You can’t perform that action at this time.
0 commit comments