[Geometry:Bugfix] Fix bug for recompute for reshape geometry

alibaba · Nov 7, 2023 · 0fb1b1e · 0fb1b1e
1 parent 91ad395
commit 0fb1b1e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/source/geometry/GeometryReshape.cpp b/source/geometry/GeometryReshape.cpp
@@ -59,7 +59,10 @@ class SingleGeometryComputer : public GeometryComputer {
         des->regions[0].origin = inputs[0];
         des->regions[0].size[0] = 1;
         des->regions[0].size[1] = 1;
-        des->regions[0].size[2] = inputs[0]->elementSize();
+        des->regions[0].size[2] = 1;
+        for (int i = 0; i < inputs[0]->dimensions(); ++i) {
+            des->regions[0].size[2] *= inputs[0]->length(i);
+        }
         des->regions[0].src.stride[2] = 1;
         des->regions[0].dst.stride[2] = 1;
         des->regions[0].src.offset = 0;