test_keras_predict_inside_function.py

import tensorflow as tf

batch_size = 32

data = tf.random.uniform((1024, 4))
model = tf.keras.Sequential((
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(2, activation='linear'),
))

output = model.predict(data, batch_size=batch_size, steps=1024//batch_size)

print(tf.shape(output))


@tf.function
def predict_model(data, batch_size):
    N = tf.shape(data)[0]
    output = model.predict(data, batch_size=batch_size, steps=N // batch_size)
    return output


output = predict_model(data, batch_size=batch_size)
print(tf.shape(output))