Error During Model Inspection with Vitis Inspector Tool

[oshaughnessya]

Issue Title:

Error During Model Inspection with Vitis Inspector Tool

Description:

When attempting to inspect my TensorFlow model using the Vitis Inspector tool, I encountered the following error:

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
])

# The training step has been removed for shortening on the report. It successfully trained.

inspector = vitis_inspect.VitisInspector(target="DPUCZDX8G_ISA1_B512")
inspector.inspect_model(model,
                        plot=True,
                        plot_file="model.svg",
                        dump_results=True,
                        dump_results_file="inspect_results.txt")

ValueError: Cannot assign value to variable 'dense_1/kernel:0': Shape mismatch. The variable shape (784, 10), and the assigned value shape (128, 10) are incompatible.

This error suggests a shape mismatch when assigning weights to the dense layer in the model being inspected.

Steps to Reproduce:

1. Define and train a basic MNIST model using the provided model definition.
2. Attempt to inspect the trained model using the Vitis Inspector tool with the provided code snippet.
3. Observe the error message mentioned above.

Expected Behavior:

I expected the model inspection process to complete without errors, providing insights into the model's structure and parameters.

Actual Behavior:

The model inspection process fails with the error message mentioned above.

Docker Image:

xilinx/vitis-ai-tensorflow2-cpu:latest

Additional Information:

I have verified that the model being inspected matches the architecture of the trained model and that there were no modifications made to the model after training.

Reproducibility:

The issue is reproducible consistently.