https://github.com/linkedin/Liger-Kernel Should be easy with https://github.com/linkedin/Liger-Kernel/issues/242#issuecomment-2347334135 Also we should patch `kl_div` loss with ligers kl_div if liger is available.
https://github.com/linkedin/Liger-Kernel
Should be easy with linkedin/Liger-Kernel#242 (comment)
Also we should patch
kl_divloss with ligers kl_div if liger is available.