MPT: Change order of operands to enable PT2 compile for inference

[tdoublep]

This is a trivial change, which doesn't actually change any logic, but it very helpful to enable PyTorch 2 compile for inference using mpt-based models.

During compilation, evaluating attention_mask[:, 0].sum() != attention_mask.shape[0] seems to cause a bunch of problems because it compares the shape of the tensor to its actual contents, and leads to graph breaks. By changing the ordering of the operands in this if statement, the problem is resolved since self.training gets evaluated first and fails, thus preventing the rest from being evaluated.

[dakinggg]

LGTM, thanks!

[tdoublep]

Anything I can do regarding the failing checks? Not 100% sure but doesn't seem to be related to the proposed code change.

[dakinggg]

Its just the autoformatting, if you run pre-commit run --all-files and commit, it will fix it.

[tdoublep]

@dakinggg I ran the code formatting, but looks like an approval is needed for the checks to re-run.