[issue 594]exposed the chunk loss in jsb_loss.py

[vulkomilev]

Summary

I have exposed the chunk loss as requested here #594

Testing Done

I have ran test/chunked_loss/test_jsd_loss.py but I couldn't finish it even with the original repo.

4090

Complete the following tasks before sending your PR, and replace [ ] with
[x] to indicate you have done them.
-->

• Hardware Type:
• [ x] run make test to ensure correctness
• [ x] run make checkstyle to ensure code style
• [x ] run make test-convergence to ensure convergence

[RichhLi]

Hey great work. Sorry I didn't see this, I worked on the same issue. #599