def fused_fwd_bwd(input_chunk, target_chunk, ref_input_chunk, chosen_nll_target_chunk):
            """
            Fused forward and backward pass for a chunk of input and target.
            """
            if bias is not None:
                return torch.func.grad_and_value(compute_loss, argnums=(0, 1, 3), has_aux=True)(
                    input_chunk,
                    weight,
                    target_chunk,
                    bias,
                    ref_input_chunk=ref_input_chunk,
                    chosen_nll_target_chunk=chosen_nll_target_chunk,
                )
            else:
                return torch.func.grad_and_value(compute_loss, argnums=(0, 1), has_aux=True)(
                    input_chunk,
                    weight,
                    target_chunk,
                    ref_input_chunk=ref_input_chunk,
                    chosen_nll_target_chunk=chosen_nll_target_chunk,
                )

            Fused forward and backward pass for a chunk of input and target.
            
src/liger_kernel/chunked_loss/fused_linear_preference.py