tensorflow: Tf lite interpreter.allocate_tensors() Runtime error

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 18.04):
• TensorFlow installed from (binary):
• TensorFlow version (2.4.1):
• Also tried tf-nightly

Ive used the following code to convert keras model to tflite. Conversion went successfull without any errors. Note that i used tf-nightly for the conversion as stable version of tf 2.4.1 was unable to convert

import tensorflow as tf
from tensorflow.keras.models import load_model

modelPath = "u2net_keras.h5"

tflite_model = load_model(modelPath)
converter = tf.lite.TFLiteConverter.from_keras_model(tflite_model)

# converter = tf.lite.TFLiteConverter.from_keras_model_file( 'u2netp_keras.h5')
tfmodel = converter.convert()
open ("model.tflite" , "wb") .write(tfmodel)
print("Sucess!")

Now I am using the following code to get inference from the converted model but code seems to get stuck on 'interpreter.allocate_tensors()' line

import tensorflow as tf
import PIL.Image as Image
import numpy as np

import cv2


def preprocess(img_path, dim):
    img = Image.open(img_path)
    img = img.resize(dim, Image.BILINEAR)
    img_data = np.array(img)
    img_data = np.transpose(img_data, [2, 0, 1])
    img_data = np.expand_dims(img_data, 0)
    mean_vec = np.array([0.485, 0.456, 0.406])
    stddev_vec = np.array([0.229, 0.224, 0.225])
    norm_img_data = np.zeros(img_data.shape).astype('float32')
    for i in range(img_data.shape[1]):
        norm_img_data[:,i,:,:] = (img_data[:,i,:,:]/255 - mean_vec[i]) / stddev_vec[i]
    return norm_img_data


w = 320
h = 320
dim = (w,h)
model = "U2net-tflite-models/u2net.tflite"
img_path = 'boat.jpg'
img = preprocess(img_path, dim)


# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model)

interpreter.resize_tensor_input(0, [1,3, 320, 320], strict=False)

interpreter.allocate_tensors()



#get input and output tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()


# Read the image and decode to a tensor

#Preprocess the image to required size and cast
input_shape = input_details[0]['shape']
print("input shape = ",np.expand_dims(img,0))
input_tensor= np.array(np.expand_dims(img,0))

#set the tensor to point to the input data to be inferred
input_index = interpreter.get_input_details()[0]["index"]
interpreter.set_tensor(input_index, input_tensor)
#Run the inference
interpreter.invoke()
output_details = interpreter.get_output_details()


output_data = interpreter.get_tensor(output_details[0]['index'])
results = np.squeeze(output_data)
top_k = results.argsort()
for label, idx in train_data_gen.class_indices.items():
    if top_k[idx]==1:
        print("Prediction: " ,label)

ERROR log

2021-04-15 18:34:01.240901: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2021-04-15 18:34:01.240937: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
Traceback (most recent call last):
  File "/home/faraz/Desktop/Esper_Solutions/upwrok/tfLite/U-2-Net-Keras/tflite_inference.py", line 34, in <module>
    interpreter.allocate_tensors()
  File "/home/faraz/.local/lib/python3.6/site-packages/tensorflow/lite/python/interpreter.py", line 420, in allocate_tensors
    return self._interpreter.AllocateTensors()
RuntimeError: tensorflow/lite/kernels/space_to_batch_nd.cc:67 op_context->block_shape->dims->data[0] != spatial_dims_num (3 != 2)Node number 49 (SPACE_TO_BATCH_ND) failed to prepare.


Process finished with exit code 1