diff options
| -rw-r--r-- | ethosu/vela/tflite_writer.py | 16 |
1 files changed, 7 insertions, 9 deletions
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py index 8f523c7d..77e6f953 100644 --- a/ethosu/vela/tflite_writer.py +++ b/ethosu/vela/tflite_writer.py @@ -168,17 +168,15 @@ class TFLiteSerialiser: tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix] custom_code_offset = builder.CreateString(code[len(custom_prefix) :]) else: - try: - tf_code, opt_serializer = builtin_operator_inv_map[code] - except KeyError: - print( - "Warning: Writing operation %s, which does not have a direct TensorFlow Lite mapping," - "as a custom operation" % (code,) - ) - tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix] + assert ( + code in builtin_operator_inv_map + ), "Vela does not contain a mapping to serialise {} operator to a TensorFlow Lite operator".format(code) + tf_code, opt_serializer = builtin_operator_inv_map[code] if tf_code == BuiltinOperator.CUSTOM: - assert code == "NpuOp" # Currently only support serialising NPU operators as a custom op + assert ( + code == "NpuOp" + ), "Vela only supports serialising NpuOp operators as TensorFlow Lite Custom operators" custom_code_offset = builder.CreateString("ethos-u") self.operator_code_map[code] = (idx, tf_code, opt_serializer) |