## Summary

> Please provide a summary of the changes made and the issue that is being addressed.

## Checklist:

- [ ] Adheres to the [style guidelines](https://facebookresearch.github.io/momentum/docs/developer_guide/style_guide)
- [ ] Codebase formatted by running `pixi run lint`

## Test Plan

> Please describe the tests you conducted to verify your changes.