[chunked loss] align teacher and student logit shape#634

shivam15s merged 2 commits intomainlinkedin/Liger-Kernel:mainfrom

kd_alignlinkedin/Liger-Kernel:kd_alignCopy head branch name to clipboard