Model runs without error in Tensorflow, but crashes with a segmentation fault in TFLite
Hi @m-decoster 

I have tried to reproduce the issue on `tf-nightly` and did not observe any crash or segmentation fault.

I was able to get the output without any error.

Please find the gist [here](https://colab.research.google.com/gist/pjpratik/fb385410acdcc281d3fa2471d4fb7179/61067.ipynb) and let us know if it helps.

Thanks. This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61067">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61067">No</a>

### 1. System information

- OS Platform and Distribution: Ubuntu 20.04.4 LTS
- TensorFlow installation: pip
- TensorFlow library: 2.12.0
- TFLite runtime: 2.12.0

### 2. Code

The model is exported from PyTorch using ONNX. I have not included the PyTorch code below for brevity's sake (and because it is used for an active Kaggle competition); you can download the saved Keras model [here](https://cloud.ilabt.imec.be/index.php/s/Dgpi9SQTcyc23wm). The TFLite conversion code is given below, but you can also download the TFLite model [here](https://cloud.ilabt.imec.be/index.php/s/Dgpi9SQTcyc23wm) (same link).

Below is the code to create and save the Keras model from two PyTorch models `feat_gen` and `model`, converted using ONNX:

```python
class TFInferModel(tf.Module):
    def __init__(self):
        super(TFInferModel, self).__init__()
        self.feat_gen = tf.saved_model.load("feat_gen.pb")
        self.model = tf.saved_model.load("model.pb")

        self.feat_gen.trainable = False
        self.model.trainable = False

    @tf.function(input_signature=[tf.TensorSpec(shape=[None, 126], dtype=tf.float32, name="inputs")])
    def call(self, inputs):
        output_tensors = {}

        # Add batch dimension.
        inputs = inputs[None]

        # Process using ported PyTorch model.
        features = self.feat_gen(inputs=inputs)["outputs"]
        outputs = self.model(inputs=features)["outputs"]

        # Remove batch dimension.
        outputs = outputs[0]

        output_tensors["outputs"] = outputs
        return output_tensors

tf_model = TFInferModel()
tf.saved_model.save(tf_model, "tf_model", signatures={"serving_default": tf_model.call})
```

The model can be loaded in Keras and run:

```python
model = tf.saved_model.load("tf_model")
inputs = tf.zeros((100, 126), dtype=tf.float32)
output = model.call(inputs=inputs)
```

It can also be converted to TFLite:

```python
converter = tf.lite.TFLiteConverter.from_saved_model("tf_model")

tf_lite_model = converter.convert()
output_path = "model.tflite"
with open(output_path, "wb") as f:
    f.write(tf_lite_model)
```

And finally the code for TFLite inference:

```python
interpreter = tflite.Interpreter(model_path="model.tflite")
prediction_fn = interpreter.get_signature_runner("serving_default")
inputs = np.zeros((100, 126), dtype=np.float32)
output = prediction_fn(inputs=inputs)
```

### 3. Failure after conversion

The Keras inference code runs without issue. The TFLite inference code crashes immediately with a segmentation fault (no further info is given).