Fix bug: Prevent inconsistent return value with `eos_early_stop` (#3)
Browse files- Fix bug: Prevent inconsistent return value with `eos_early_stop` (d72fd9e3d3fc7b506fd937c295e5e177a51d578f)
Co-authored-by: Yao <[email protected]>
- modeling_llada2_moe.py +1 -1
modeling_llada2_moe.py
CHANGED
|
@@ -1618,4 +1618,4 @@ class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin):
|
|
| 1618 |
first_mask_position = mask_positions[0].item()
|
| 1619 |
else:
|
| 1620 |
first_mask_position = gen_length
|
| 1621 |
-
return generated_answer[:,
|
|
|
|
| 1618 |
first_mask_position = mask_positions[0].item()
|
| 1619 |
else:
|
| 1620 |
first_mask_position = gen_length
|
| 1621 |
+
return generated_answer[:, : input_ids.shape[1] + first_mask_position + 1]
|