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

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