Browse Source

fix(mge): fix dimshuffle shape infer

GitOrigin-RevId: ad7a73fbd1
release-1.10
Megvii Engine Team 3 years ago
parent
commit
275f12c98a
1 changed files with 1 additions and 0 deletions
  1. +1
    -0
      imperative/src/impl/ops/specializations.cpp

+ 1
- 0
imperative/src/impl/ops/specializations.cpp View File

@@ -81,6 +81,7 @@ std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible(
src.layout.ndim);
size_t idx = 0;
bool input_used[TensorLayout::MAX_NDIM] = {0};
out_shape.ndim = ds.pattern.size();
for (auto i : ds.pattern) {
if (i < 0) {
out_shape[idx] = 1;


Loading…
Cancel
Save