Fix failing tests on main due to torch 2.1

[ydshieh]

What does this PR do?

Fix failing tests on main due to torch 2.1

[ydshieh]

for Hubert and Wav2Vec2, fx tracing starts to fail with torch 2.1

[ydshieh]

despite the changes in the modeling files, this part still fails

[ArthurZucker]

Thanks, let's try to use unitest skip and put the torch fx proxy in the model forward

[ArthurZucker]

Instead of commenting we can add something like

self.skipTest("Skipping until we fix it ")

[ArthurZucker]

Same here

[ArthurZucker]

let's do the check a lot earlier, we can check this in the model's forward only once!

[ydshieh]

I am not sure about this place. This block contain

            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

and it should be run no matter if we are in fx proxy or not

[LysandreJik]

LGTM, let's ensure we check for eventual performance issues and fix the TODOs quickly.

Thanks for fixing these tests so quickly!

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[ArthurZucker]

Thanks

[ydshieh]

@michaelbenayoun Could you help us on the torch fx tests for wav2vec2/hubert with torch 2.1.

See this internal discussion

But in short, it can't do

if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len)

as the corresponding proxy object has no _metadata attribute (with torch 2.1) but it has it with torch 2.0.