Getting Value Error while running the model prediction when it is deployed to Google Cloud Run

I’m getting the following error while running the model prediction through my streamlit app deployed on Google Cloud Run. But the same code is running perfectly locally and on Google colab with same python version(3.10). The only difference is tensorflow version the deployed version using 2.14 whereas on the local machine I’m using 2.15.

But I don’t see major changes in both versions of tensorflow that would affect the normal prediction code.

ValueError: Cannot reshape a tensor with 3962112 elements to shape [67,77,16,16] (1320704 elements) for '{{node Reshape_2}} = Reshape[T=DT_INT64, Tshape=DT_INT32](ArgMax, Reshape_2/shape)' with input shapes: [15477,16,16], [4] and with input tensors computed as partial shapes: input[1] = [67,77,16,16].
Traceback:
File "/home/otbuser/.local/lib/python3.10/site-packages/streamlit/runtime/scriptrunner/exec_code.py", line 85, in exec_func_with_error_handling
    result = func()
File "/home/otbuser/.local/lib/python3.10/site-packages/streamlit/runtime/scriptrunner/script_runner.py", line 576, in code_to_exec
    exec(code, module.__dict__)
File "/app/pages/9-Performing-like-an-algorithm.py", line 419, in <module>
    a_r, predicted_array = prediction(imagepath+inputimage, modelpath+model)
File "/app/otbtf/unet_helper.py", line 208, in prediction
    segmentation_patches = tf.reshape(
File "/opt/otbtf/lib/python3/dist-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
File "/opt/otbtf/lib/python3/dist-packages/tensorflow/python/framework/ops.py", line 1021, in _create_c_op
    raise ValueError(e.message)

Dockerfile I used to deploy my streamlit application

FROM mdl4eo/otbtf:4.3.1-cpu


WORKDIR /app

# Copy your application code to the container
COPY . ./

# CMD ls -l /home/otbuser/
# Install Python dependencies including otbtf
RUN pip install --no-cache-dir -r requirements.txt

# port
EXPOSE 8080

ENTRYPOINT ["streamlit", "run", "Planetary-Computing-Logics.py", "--server.port=8080", "--server.address=0.0.0.0"]


The the particular docker image is used as I need it for other models. Tensorflow is provided by this base image and I can’t update the tensorflow version.

Please suggest ways to resolve the above value error.