Bug fix for missing distillation loss arguments. (#983)

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Bug fix for missing distillation loss arguments.

<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Bug fix for missing distillation loss arguments.

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

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

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [] run `make test-convergence` to ensure convergence

src/liger_kernel/chunked_loss/cosine_similarity_loss.py

@@ -9,7 +9,13 @@ from liger_kernel.chunked_loss.fused_linear_distillation import LigerFusedLinear
 
 class LigerFusedLinearCosineSimilarityFunction(LigerFusedLinearDistillationBase):
     @staticmethod
-    def distillation_loss_fn(student_logits, teacher_logits, beta=1.0):
+    def distillation_loss_fn(
+        student_logits,
+        teacher_logits,
+        target=None,
+        ignore_index=None,
+        beta=1.0,
+    ):
         """
         Compute Cosine loss (Cosine Similarity Loss).
         Args:

src/liger_kernel/chunked_loss/fused_linear_distillation.py

@@ -13,6 +13,8 @@ class LigerFusedLinearDistillationBase(torch.autograd.Function):
     def distillation_loss_fn(
         student_logits,
         teacher_logits,
+        target=None,
+        ignore_index=None,
     ):
         """
         Compute distillation loss.

test/chunked_loss/test_cosine_loss.py

@@ -33,7 +33,7 @@ class HFCosineLoss(HFDistillationLoss):
             temperature=temperature,
         )
 
-    def distillation_loss(self, student_logits, teacher_logits, beta=1.0):
+    def distillation_loss(self, student_logits, teacher_logits, target=None, ignore_index=None, beta=1.0, **kwargs):
         # Compute normalized logits
         print(f"student_logits.shape: {student_logits.shape}")
         student_norm = F.normalize(student_logits, p=2, dim=-1)