**Suggested improvement:** It would be good to add a "division by zero" check in chunked loss handling in [kd_losses.py](https://github.com/pytorch/torchtune/blob/d23fa93d21045ab88f6914051f11de2b9afcd9eb/torchtune/modules/loss/kd_losses.py#L140). **Context:** This is based on the discussion in PR #2094. **Potential issue:** [ForwardKLWithChunkedOutputLoss](https://github.com/pytorch/torchtune/blob/d23fa93d21045ab88f6914051f11de2b9afcd9eb/torchtune/modules/loss/kd_losses.py#L140) does not check for division by zero, while the non-chunked version does.