#0: Fix for rank 0 and 1 tensors

tenstorrent · Jan 16, 2025 · c261c8c · c261c8c
1 parent db1d0cb
commit c261c8c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/ttnn/cpp/ttnn/tensor/layout/tensor_layout.cpp b/ttnn/cpp/ttnn/tensor/layout/tensor_layout.cpp
@@ -262,7 +262,10 @@ Size TensorLayout::get_physical_shard_shape() const {
 
 Size TensorLayout::compute_logical_2d_shape(const ttnn::SimpleShape& shape) const {
     size_t width = shape[-1];
-    size_t height = shape.volume() / width;
+    size_t height = 1;
+    for (int i = -2; i >= -shape.rank(); --i) {
+        height *= shape[i];
+    }
     return Size{height, width};
 }