Address comments

Author: NicolaLancellotti

python/tvm/contrib/ethosu/cascader/device_config.py

@@ -55,9 +55,9 @@ def __init__(self, shape: List[int], layout="NHWC"):
                 self.width = int(shape[2])
                 self.depth = int(shape[3])
             elif length == 3:
-                self.height = int(shape[1])
-                self.width = int(shape[2])
-                self.depth = 1
+                self.height = int(shape[0])
+                self.width = int(shape[1])
+                self.depth = int(shape[2])
             elif length == 2:
                 self.height = int(shape[0])
                 self.width = int(shape[1])

python/tvm/relay/backend/contrib/ethosu/te/identity.py

@@ -131,8 +131,8 @@ def match_ethosu_identity(output_tensor, device_config):
     ofm_dtype = output_tensor.dtype
 
     input_tensors_shape = input_tensors[0].shape
-    ifm_channels = int(input_tensors_shape[3] if len(input_tensors_shape) > 3 else 1)
-    ofm_channels = ifm_channels
+    length = len(input_tensors_shape)
+    channels = int(input_tensors_shape[length - 1]) if length >= 3 else 1
 
     subkernels = len(device_config.get_kernel_steps(identity.op.name, 1, 1, ifm_dtype))
 
@@ -143,8 +143,8 @@ def match_ethosu_identity(output_tensor, device_config):
         propagators[0],
         identity.op.attrs,
         output_tensor.shape,
-        ofm_channels,
-        ifm_channels,
+        channels,
+        channels,
         output_layout,
         input_layout,
         ifm_dtype,