Tensorflow custom layer rbf_for_tf2 to ONNX

Question

I'm trying to convert a model with tf2onnx that is using a custom RBF layer in Tensorflow 2. It throws me this error and I'm not sure if it's related to tf2onnx or the implementation of the layer. Anyone had a similar issue?

TF 2.12.1 tf2onnx: 0.4.1

RBF Layer used

Error:

Failed to apply optimize_transpose
Traceback (most recent call last):
  File "D:\Development\TensorflowProjects\venv\lib\site-packages\tf2onnx\optimizer\__init__.py", line 62, in optimize_graph
    graph = opt.optimize(current, iteration) or graph
  File "D:\Development\TensorflowProjects\venv\lib\site-packages\tf2onnx\optimizer\optimizer_base.py", line 41, in optimize
    graph = self._optimize(graph)
  File "D:\Development\TensorflowProjects\venv\lib\site-packages\tf2onnx\optimizer\transpose_optimizer.py", line 158, in _optimize
    return self._apply_optimization(graph, self._optimize_at_current_graph_level)
  File "D:\Development\TensorflowProjects\venv\lib\site-packages\tf2onnx\optimizer\optimizer_base.py", line 62, in _apply_optimization
    graph = optimize_func(graph)
  File "D:\Development\TensorflowProjects\venv\lib\site-packages\tf2onnx\optimizer\transpose_optimizer.py", line 192, in _optimize_at_current_graph_level
    self.post_optimize_action()
  File "D:\Development\TensorflowProjects\venv\lib\site-packages\tf2onnx\optimizer\transpose_optimizer.py", line 124, in post_optimize_action
    new_shape = _calculate_new_shape(self._g, op)
  File "D:\Development\TensorflowProjects\venv\lib\site-packages\tf2onnx\optimizer\transpose_optimizer.py", line 99, in _calculate_new_shape
    perm_shape = [tagged_shape[p] for p in perm]
  File "D:\Development\TensorflowProjects\venv\lib\site-packages\tf2onnx\optimizer\transpose_optimizer.py", line 99, in <listcomp>
    perm_shape = [tagged_shape[p] for p in perm]
IndexError: list index out of range

Conversion:

# ONNX Save
input_signatures = [tf.TensorSpec([1, 3], tf.float32, name='input')]
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signatures, opset=13)
onnx.save(onnx_model, 'trained_model.onnx')