[CI / core] Fix CI with GC + pytorch 2.2

[younesbelkada]

What does this PR do?

Fixes the current failing CI for Mistral, Mixtral and Qwen2 for gradient checkpointing. For some reason, since pytorch 2.2, gradient checkpointing raises an error when going through in-place operations such as tensor.mul_(xxx) which was not the case in earlier versions.

Simply replacing causal_mask.mul_(~torch.eq(causal_mask, causal_mask.min()).all(dim=-1)[..., None]) by causal_mask = causal_mask * (~torch.eq(causal_mask, causal_mask.min()).all(dim=-1)[..., None])

This makes me think we should maybe have a job that runs the CI on torch nightly to catch these early bugs, do we have that already? If not, happy to have a look

cc @ArthurZucker @amyeroberts

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[younesbelkada]

It appears the rootcause was slightlly different, see: #29027