Cannot compile model for inferentia: ValueError: Attempt to convert a value (None) with an unsupported type (<class 'NoneType'>) to a Tensor.

Currently, I host my model with tensorflow_model_server. Here is how I export my model:

model = tf.keras.models.load_model("model.hdf5")

def __decode_images(images, nch):
    o = tf.vectorized_map(lambda x: tf.image.decode_jpeg(x, nch), images)
    o = tf.image.resize(o, (128,128))
    o = tf.cast(o, dtype=tf.float16) / 255
    o = tf.reverse(o, axis=[-1])  # RGB2BGR
    return o


def __encode_images(images):
    images = tf.image.convert_image_dtype(images, tf.uint8, saturate=True)
    o = tf.vectorized_map(tf.image.encode_jpeg, images)
    return o


@tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string, name='image')])
def serving(img):
    img = __decode_images(img, 3)
    o = model([img], training=False)
    o = __encode_images(o)
    return {
        'output': o,
    }

tf.saved_model.save(model, export_dir=args.output, signatures=serving

Now, I want to compile my model for inferentia. Here is that I come up with:

mage_path = '1.jpg' 
with tf.io.gfile.GFile(image_path, 'rb') as f:
    image_data = f.read()
sample_input = tf.constant([image_data])
model_neuron = tfn.trace(serving, sample_input)
model_neuron.save(args.output)

But I get the following error:

    model_neuron = tfn.trace(serving, sample_input)
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow_neuron/python/_trace.py", line 215, in trace
    graph_def = gdu.inline_shape_inputs_in_subgraphs(graph_def)
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow_neuron/python/graph_def_util.py", line 245, in inline_shape_inputs_in_subgraphs
    shape_tensor = convert_to_tensor(tensor_content, dtype)
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow/python/profiler/trace.py", line 183, in wrapped
    return func(*args, **kwargs)
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 1638, in convert_to_tensor
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py", line 343, in _constant_tensor_conversion_function
    return constant(v, dtype=dtype, name=name)
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py", line 267, in constant
    return _constant_impl(value, dtype, shape, name, verify_shape=False,
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py", line 279, in _constant_impl
    return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py", line 304, in _constant_eager_impl
    t = convert_to_eager_tensor(value, ctx, dtype)
  File "/home/mosi/miniconda3/envs/inf/lib/python3.8/site-packages/tensorflow/python/framework/constant_op.py", line 102, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
ValueError: Attempt to convert a value (None) with an unsupported type (<class 'NoneType'>) to a Tensor.