TODO write me Relevant links: * https://github.com/pytorch/torchtitan/pull/808 * https://github.com/pytorch/ao/pull/1718