RESHAPE failed to prepare when invoking tf.signal.stft in tflite

I am building a flutter app that needs to record an audio and predict some label using a tflite model I built. For linking the audio recording and tflite I use the flutter plugin tf-lite audio (https://github.com/Caldarie/flutter_tflite_audio).

The tensorflow model works on colab but when I launch the app and inference happens hence when it calls interpreter.invoke(), the following error occurs:

TensorFlow Lite Error: tensorflow/lite/kernels/reshape.cc:58 stretch_dim != -1 (0 != -1)
TensorFlow Lite Error: Node number 26 (RESHAPE) failed to prepare.
Failed to invoke the interpreter with error: Must call allocateTensors().
2
Fatal error: Unexpectedly found nil while implicitly unwrapping an Optional value: file tflite_audio/SwiftTfliteAudioPlugin.swift, line 290
* thread #2, queue = 'conversionQueue', stop reason = Fatal error: Unexpectedly found nil while implicitly unwrapping an Optional value
    frame #0: 0x00000001a672ee08 libswiftCore.dylib`_swift_runtime_on_report
libswiftCore.dylib`_swift_runtime_on_report:
->  0x1a672ee08 <+0>: ret
libswiftCore.dylib`_swift_reportToDebugger:
    0x1a672ee0c <+0>: b      0x1a672ee08               ; _swift_runtime_on_report
libswiftCore.dylib`_swift_shouldReportFatalErrorsToDebugger:
    0x1a672ee10 <+0>: adrp   x8, 341475
    0x1a672ee14 <+4>: ldrb   w0, [x8, #0x7c8]
Target 0: (Runner) stopped.
Lost connection to device.

This error message appears even though I added allocateTensors in the SwiftTfliteAudioPlugin.swift file here:

var interval: TimeInterval!
        var outputTensor: Tensor!
        
        do {
            // Copy the `[Int16]` buffer data as an array of Floats to the audio buffer input Tensor.
            let audioBufferData = Data(copyingBufferOf: buffer.map { Float($0) / maxInt16AsFloat32 })
            try interpreter.copy(audioBufferData, toInputAt: 0)
            
            // I added this line
            try interpreter.allocateTensors()
            
            // Calculate inference time
            let startDate = Date()
            try interpreter.invoke() //required!!! Do not touch
            interval = Date().timeIntervalSince(startDate) * 1000
            
            // Get the output `Tensor` to process the inference results.
            outputTensor = try interpreter.output(at: 0)
            print(outputTensor as Any)

        } catch let error {
            print("Failed to invoke the interpreter with error: \(error.localizedDescription)")
        }

In the tflite model here is the problematic node on netron

It looks like it is only squeezing the first dimension so maybe it cannot because as you can see on the summary of my model the first dimension is None, I tried some tricks to avoid having this None but I am not familiar enough with tensorflow to be sure about the validity of the operations I am doing.

I have boiled down my model to the minimal size and this node is between these 2 lines of code, so I suspect the tf.signal.stft function to do this reshaping but have no idea.

spectrograms = tf.signal.stft(waveforms,
                                     frame_length=self.fft_size,
                                     frame_step=self.hop_size,
                                     pad_end=False)

magnitude_spectrograms = tf.abs(spectrograms)

Can anyone help on this issue? Thanks!