ONNX runtime cannot interpret valid TF2 model

Question

I'm trying to get a Tensorflow2 model to run in ONNX runtime, but it is failing when I do custom computation at the end of the model. I think it might be something to do with the way I index my variables. Are there any ONNX experts who can help?

The only big hint is that the code between here! and to here is responsible for the error. If I remove it and use a simple sigmoid output, then the whole thing works.

fd = np.random.rand(100,16)
fl = np.random.randint(0,1,[100,])
D = 0.4375
def get_m():
    I = Input(shape=(None, 16))
    d1 = Dense(16, kernel_regularizer='l2', activation='tanh')(I)
    dr1 = Dropout(D)(d1)
    d2 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr1)
    dr2 = Dropout(D)(d2)
    d3 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr2)
    dr3_1 = Dropout(D)(d3)
    dr3_2 = Dropout(D)(d3)
    d4_1 = Dense(1, kernel_regularizer='l2', activation='relu')(dr3_1)
    d4_2 =  Dense(1, kernel_regularizer='l2', activation='relu')(dr3_2)
    # here!
    c = Concatenate()([d4_1, d4_2, I])
    def custom_layer(x):
        return tf.exp(-x[:, 4] / (1e-9 + x[:, 0]) - x[:, 5] / (1e-9 + x[:, 1]))

    hl = tf.keras.layers.Lambda(custom_layer, output_shape=(1,))(c)
    #d5 = Dense(1, activation='sigmoid')(c)
    # to here
    return Model(inputs=I, outputs=hl)

model = get_m()
model.compile(loss=tf.keras.losses.BinaryFocalCrossentropy(), optimizer=tf.keras.optimizers.Adam(), metrics=[tf.keras.metrics.AUC(), tf.keras.metrics.Precision(), tf.keras.metrics.Recall()])
model.fit(x=fd, y=fl, epochs=100, validation_split=0.1
tf2onnx.convert.from_keras(model, output_file='tf2_neural_model_2deep.onnx')

And here's the code I use to run it:

ort = require('onnxruntime-node')
sess = await ort.InferenceSession.create('tf2_neural_model_2deep.onnx')
rowVectors = [[0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]]
inputTensor = new ort.Tensor('float32', rowVectors.flat(), [rowVectors.length, 1, rowVectors[0].length]);
output = await sess.run({'input_38': inputTensor});

It gives me this error:

2023-09-12 09:25:00.950381910 [E:onnxruntime:, sequential_executor.cc:514 ExecuteKernel] 
Non-zero status code returned while running Squeeze node. 
Name:'model_46/lambda_41/strided_slice__288' Status Message: 
/onnxruntime_src/onnxruntime/core/providers/cpu/tensor/squeeze.h:52 static
onnxruntime::TensorShapeVector onnxruntime::SqueezeBase::ComputeOutputShape(const 
onnxruntime::TensorShape&, const TensorShapeVector&) input_shape[i] == 1 was false. 
Dimension of input 1 must be 1 instead of 0. shape={1,0,18}

Uncaught:
Error: Non-zero status code returned while running Squeeze node. 
Name:'model_46/lambda_41/strided_slice__288' Status Message: 
/onnxruntime_src/onnxruntime/core/providers/cpu/tensor/squeeze.h:52 static 
onnxruntime::TensorShapeVector onnxruntime::SqueezeBase::ComputeOutputShape(const 
onnxruntime::TensorShape&, const TensorShapeVector&) input_shape[i] == 1 was false. 
Dimension of input 1 must be 1 instead of 0. shape={1,0,18}