Fix E731: Convert lambda assignments to def functions

[Copilot]

Reduces ignored ruff rules by fixing all E731 violations ("do not assign a lambda expression, use a def").

Changes

• apex/contrib/openfold_triton/layer_norm.py: Convert 3 grid lambda assignments to def functions
• apex/contrib/openfold_triton/mha.py: Convert 2 grid lambda assignments to def functions
• examples/imagenet/main_amp.py: Convert collate_fn lambda to def function
• pyproject.toml: Remove E731 from ignore list

Example

# Before
grid = lambda kwargs: (triton.cdiv(kwargs["M"], kwargs["M_BLOCK"]),)

# After
def grid(kwargs):
    return (triton.cdiv(kwargs["M"], kwargs["M_BLOCK"]),)

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Reduce the number of ignored rules</issue_title> <issue_description>Create a pull request per ignored rule listed in https://github.com/NVIDIA/apex/blob/87868b0e3a39c5f4364ac2d8131a7e63fce0f1a4/pyproject.toml#L8-L25</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes NVIDIA/apex#1964

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.

[crcrpar]

pre-commit.ci autofix