diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py index 4eaa8280c307..05b781b62300 100644 --- a/python/tvm/relay/frontend/tflite.py +++ b/python/tvm/relay/frontend/tflite.py @@ -1300,11 +1300,6 @@ def convert_square(self, op): assert len(output_tensors) == 1, "output tensors length should be 1" output_tensor = output_tensors[0] - if self.is_quantized(op): - raise tvm.error.OpNotImplemented( - "TFlite quantized SQUARE operator is not supported yet." - ) - exp_type = self.get_tensor_type_str(output_tensor.tensor.Type()) out = _op.power(in_expr, relay.const(2, exp_type)) diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py index f1e2e28f086b..bade0c66e9cf 100644 --- a/tests/python/frontend/tflite/test_forward.py +++ b/tests/python/frontend/tflite/test_forward.py @@ -2414,7 +2414,7 @@ def test_all_unary_elemwise(): _test_forward_unary_elemwise(_test_floor) _test_forward_unary_elemwise(_test_exp) _test_forward_unary_elemwise(_test_log, negative=False) - _test_forward_unary_elemwise(_test_square) + _test_forward_unary_elemwise(_test_square, int_quant_dtype=tf.int8) _test_forward_unary_elemwise(_test_sin) _test_forward_unary_elemwise(_test_neg) _test_forward_unary_elemwise(_test_sqrt, negative=False)