From 275f12c98abb5bcbd5633f24f2b26ad8cd69f416 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Tue, 12 Apr 2022 13:32:12 +0800
Subject: [PATCH] fix(mge): fix dimshuffle shape infer

GitOrigin-RevId: ad7a73fbd1670aad926772d58b6e31fe52eeb448
---
 imperative/src/impl/ops/specializations.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/imperative/src/impl/ops/specializations.cpp b/imperative/src/impl/ops/specializations.cpp
index ff82a8407..a3e1606cf 100644
--- a/imperative/src/impl/ops/specializations.cpp
+++ b/imperative/src/impl/ops/specializations.cpp
@@ -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;
-- 
GitLab