Fix dtype mismatch in fused_linear_cross_entropy_forward

Fixes linkedin#305

Fix dtype mismatch in fused_linear_cross_entropy_forward function.

* Cast `logits_chunk` to the data type of `_input_chunk` before performing operations on it.

---

For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/linkedin/Liger-Kernel/issues/305?shareId=XXXX-XXXX-XXXX-XXXX).

src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -70,7 +70,7 @@ def fused_linear_cross_entropy_forward(
         n_non_ignore = (target_chunk != ignore_index).sum().item()
 
         # when doing CE, use the upcasted precision
-        logits_chunk = logits_chunk.float()
+        logits_chunk = logits_chunk.to(_input_chunk.dtype)
 
         # ensure _input and target are contiguous
         logits_chunk = logits_chunk.contiguous()