Refactor pixtal/loader.py and fix load_inputs function

[devisettymahidhar608]

Ticket

Link to Github Issue

Problem description

Debug PCC drop in the model

What's changed

Upon inspecting the load_inputs function, it returns the following keys: input_ids and attention_mask.

While running the model from Hugging Face in Google Colab, I observed a difference in the input keys being passed to the model.

The following keys are being sent to the model: ['input_ids', 'attention_mask', 'pixel_values', 'image_sizes'].

loader.py only sends input_ids and attention_mask; pixel_values and image_sizes are not included.

Edited loader.py to ensure correct inputs are passed to the model.

While running the Mistral/pixtral model, I encountered the following error: loc("set-dimension-size.60"): Shardy propagation only supports ranked tensors with a static shape. The tensor type causing the issue is tensor>. pixtral.log

Issue originats from the masked_scatter - sanity.log

Checklist

• New/Existing tests provide coverage for changes

[kmabeeTT]

Adding @AleksKnezevic - Thanks for changes @devisettymahidhar608 - I'd like to get some input from Aleks on what to do here, since this model is currently running (but with low PCC lately, and incorrect loader.py), do we merge these changes which causes the model to not run e2e anymore until the issue you opened tenstorrent/tt-xla#2924 gets assigned and debugged.

Side Note - We should hold off a bit on merging this until tt-forge-models recent uplift issues (a week out of date in tt-xla) get resolved.

[AleksKnezevic]

Commented on the other issue, we can't support any dynamic shapes through the compiler right now. We'll need to debug why we're seeing this dynamism.