Tape.gradient returns None on streamlit Cloud, but works fine locally

Hello,

please, does anyone have clue why grad.tape does not work speciffically on streamlit? It returns None. The arguments class_output and activations look exactly the same as on my machine, so I am clueless. Here is the problematic part of the code:

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
    st.session_state.activations, predictions = grad_model(img_tensor)
    
    # Debug: Print shape of predictions and activations
    print(f"Predictions shape: {predictions.shape}")
    print(f"Activations shape: {st.session_state.activations[0].shape}")

    print(f"Activations1: {st.session_state.activations[0]}")

    
    class_idx = np.argmax(predictions[0])  # Get the class index of the highest prediction
    print(f"prediction h1gh class idx: {class_idx}")
    st.session_state.class_output = predictions[0][class_idx]  # Access the output corresponding to that class
    print(f"prediction: {st.session_state.class_output}")

# Compute the gradient of the class output w.r.t. the activations
grads = tape.gradient(st.session_state.class_output, st.session_state.activations)

Thank you!