Grad cam works on my machine but not on streamlit cloud

Dear Streamlit experts, please help me. Locally everythink works fine to me, but on streamlit i get following error in my Keras gradcam:

File "/mount/src/applepear/src/grad_cam.py", line 20, in grad_cam
    activations, predictions = grad_model(img_tensor)
                               ^^^^^^^^^^^^^^^^^^^^^^File "/home/adminuser/venv/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from NoneFile "/home/adminuser/venv/lib/python3.12/site-packages/keras/src/ops/function.py", line 179, in _run_through_graph
    output_tensors.append(tensor_dict[id(x)])

This is the part of the code, where the problem appears:

def grad_cam(model, img_tensor, layer_name):
inputs = tf.keras.Input(shape=(28, 28, 1))
outputs = model(inputs)
# Create a model that gives us both the activations and predictions
# Ensure the model’s last convolutional layer is passed
grad_model = tf.keras.models.Model(
inputs=[model.inputs],
outputs=[model.get_layer(layer_name).output, model.output]
)

with tf.GradientTape() as tape:
    tape.watch(img_tensor)  # Ensure the input image tensor is being watched
    
    # Get the activations and predictions from the model
    activations, predictions = grad_model(img_tensor)

Also I want to mention, that the model works fine on streamlit cloud, just the gradcam shows error

Thank you for any tips

Katerina

Okay I tried this:
grad_model = tf.keras.models.Model(
#inputs=[model.inputs],
inputs=[st.session_state.model.get_layer(index=0).input],
#outputs=[model.get_layer(layer_name).output, model.output]
outputs=[st.session_state.model.layers[-4].output, st.session_state.model.layers[-1].output]
)

and now the input and shapes are fine, but there is another error, the resolving gradients are None. Somehow the resulted activations form the model dissapear in streamlit. Why?

Okay, I managed to solve this but now i have problem, that:

grads = tape.gradient(st.session_state.class_output, st.session_state.activations)

is none, even thoug the arguments seem reasonable