diff --git a/mace/python/tools/tf_converter_lib.py b/mace/python/tools/tf_converter_lib.py
index dc989796be2b852836b9f55936f251d69e8c637a..7177a691dabae5bac8fe0fd884d05850d4bac586 100644
--- a/mace/python/tools/tf_converter_lib.py
+++ b/mace/python/tools/tf_converter_lib.py
@@ -786,10 +786,9 @@ class TFConverter(object):
     reshape_op = self.tf_graph[softmax_op.name][0]
     self.unused_tensor.add(get_input_tensor(reshape_op, 1).name)
 
-    if reshape_op.outputs[0].shape.ndims == 2:
-      shape = [dim.value for dim in reshape_op.outputs[0].shape]
-      if len(shape) == 2:
-        shape = [1, 1, shape[0], shape[1]]
+    shape = [dim.value for dim in reshape_op.outputs[0].shape]
+    if len(shape) == 2:
+      shape = [1, 1, shape[0], shape[1]]
     op_def.output.extend([output.name for output in reshape_op.outputs])
     self.add_output_shape([shape], op_def)
     self.resolved_ops[reshape_op.name] = 1