[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;