Enable torch.compile for loss_fn

[BoyuanFeng]

This PR enables torch.compile for loss functions, which mitigates the performance gap between PyTorch and JAX for fastmri (11.6% -> 1.83%), librispeech_deepspeech (12.02% -> -15.01%), ogbg (6.06% -> 3.28%), and wmt (19.75% -> 3.52%).

Close: #467
Close: #488

[github-actions]

MLCommons CLA bot:
Thank you very much for your submission, we really appreciate it. Before we can accept your contribution, we ask that you sign the MLCommons CLA (Apache 2). Please use this [Google form] (https://forms.gle/Ew1KkBVpyeJDuRw67) to initiate authorization. If you are from an MLCommons member organization, we will request that you be added to the CLA. If you are not from a member organization, we will email you a CLA to sign. For any questions, please contact [email protected].
0 out of 1 committers have signed the MLCommons CLA.
❌ @boyuan Feng
Boyuan Feng seems not to be a GitHub user. You need a GitHub account after you become MLCommons member. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You can retrigger this bot by commenting recheck in this Pull Request}

[priyakasimbeg]

Thanks for getting this PR in!
It seems like there is a small linting error from one of the lines in submission_runner.py. Could you fix this please? If you have yapf==0.32 installed you can just run yapf -i -r -vv -p submission_runner.py and commit the changes.