diff --git a/imperative/src/impl/ops/specializations.cpp b/imperative/src/impl/ops/specializations.cpp
index ff82a840764643b3b6a7600a3285bbac89911ab6..a3e1606cff736388899883a89fc65db5a6d376ff 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;