Fix: pytree warning

[i-colbert]

Reason for this PR

Warning from torch 2.4.0+rocm6.1

FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.
  _register_pytree_node(immutable_dict, _immutable_dict_flatten, _immutable_dict_unflatten)

Changes Made in this PR

Using torch.utils._pytree.register_pytree_node instead.

Testing Summary

Risk Highlight

• This PR includes code from another work (please detail).
• This PR contains API-breaking changes.
• This PR depends on work in another PR (please provide links/details).
• This PR introduces new dependencies (please detail).
• There are coverage gaps not covered by tests.
• Documentation updates required in subsequent PR.

Checklist

• Code comments added to any hard-to-understand areas, if applicable.
• Changes generate no new warnings.
• Updated any relevant tests, if applicable.
• No conflicts with destination dev branch.
• I reviewed my own code changes.
• Initial CI/CD passing.
• 1+ reviews given, and any review issues addressed and approved.
• Post-review full CI/CD passing.