Export tensorflow model to ONNX and specify variable names

I have a tensorflow model written through model subclassing and I want to export it to ONNX format. This is simple enough with the script attached. However, the name of the input variable to the ONNX model is args_0. How can I rename it?

import tensorflow as tf
import tf2onnx
from tensorflow.python.keras import Model
from tensorflow.python.keras.layers import Dense


class MyModel(Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dense = Dense(16)

    def call(self, inputs, **kwargs):
        return self.dense(inputs)

    def to_onnx(self, output_path, opset=14):
        model_proto, _ = tf2onnx.convert.from_keras(
            self,
            input_signature=[tf.TensorSpec((1, 128))],
            opset=opset,
            output_path=output_path,
        )
        return


if __name__ == "__main__":
    output_path = "./test.onnx"
    A = MyModel()
    A.to_onnx(output_path)

you can provide the input name in input_signature as name="input_name" so, it should look like

 input_signature=[tf.TensorSpec((1, 128), name="input_name")],

as shown in this notebook