#4003: use if-else instead of try-except in ttnn.reshape and ttnn.per…

…mute
tenstorrent · Dec 7, 2023 · 44443f0 · 44443f0
1 parent a18444f
commit 44443f0
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/ttnn/core.py b/ttnn/core.py
@@ -1009,10 +1009,10 @@ def ttnn_reshape(ttl_input_tensor, shape):
                 ttl_input_tensor, shape
             )
 
-    try:
+    if len(input_tensor.shape) == 4 and len(shape) == 4:
         w, z, y, x = shape
         return Tensor(ttl.tensor.reshape(ttl_input_tensor, w, z, y, x))
-    except:
+    else:
 
         def torch_reshape(tensor, shape):
             return tensor.reshape(shape).contiguous()
@@ -1063,9 +1063,9 @@ def permute(input_tensor: Tensor, order: Tuple[int, ...]) -> Tensor:
 
     ttl_input_tensor = input_tensor._tensor
 
-    try:
+    if len(ttl_input_tensor.shape) == 4:
         return Tensor(ttl.tensor.permute(ttl_input_tensor, order))
-    except:
+    else:
 
         def torch_permute(tensor, order):
             return tensor.permute(order).contiguous()