diff options
Diffstat (limited to 'ethosu/vela/operation.py')
| -rw-r--r-- | ethosu/vela/operation.py | 6 |
1 files changed, 2 insertions, 4 deletions
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py index 8189793e..4a56f1f0 100644 --- a/ethosu/vela/operation.py +++ b/ethosu/vela/operation.py @@ -879,10 +879,8 @@ class Operation: else: # Special case, handled in graph optimization self.ifm_shapes.append(Shape4D(ifm_tensor.get_full_shape())) - if len(self.ofm.shape) == 2: - self.ofm_shapes.append(Shape4D([self.ofm.shape[0], 1, 1, self.ofm.shape[1]])) - else: - self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape())) + self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape())) + elif self.type == Op.Softmax: self.ifm_shapes.append(Shape4D(ifm_tensor.get_full_shape())) self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape())) |