[Fix] Register ModelOutputs as a PyTree node to enable FSDP pre_backward_hook injection

[HAOCHENYE]

• fix bug
• [Fix] Register ModelOutputs as a PyTree node to enable FSDP pre_backward_hook injection

[HAOCHENYE]

@claude fix the lint

[claude]

Claude: Fixed the lint errors as requested.

What was requested: Fix lint failures in CI.

Issues found (from CI logs):

1. ruff-format: List comprehension in _flatten_pydantic_model_with_keys needed to be collapsed to a single line.
2. mypy: Return type mismatch at line 280 — GetAttrKey not recognized as subtype of KeyEntry without explicit annotation.

Fix applied:

• Added explicit type annotation list[tuple[_pytree.KeyEntry, Any]] on key_children to satisfy mypy.
• Collapsed the list comprehension element to a single line to satisfy ruff-format.

Commit: HAOCHENYE@e653fbb