Fix TFLite RESHAPE assert

apache · Nov 12, 2019 · 4dc155e · 4dc155e
1 parent 03a29da
commit 4dc155e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
@@ -265,7 +265,7 @@ def convert_reshape(self, op):
 
         assert isinstance(op, Operator)
         input_tensors = self.get_input_tensors(op)
-        assert len(input_tensors) == 2, "input tensors length should be 2"
+        assert len(input_tensors) in [1, 2], "input tensors length should be 1 or 2"
         input_tensor = input_tensors[0]
         input_tensor_idx = input_tensor.tensor_idx