Fix bug: Prevent inconsistent return value with `eos_early_stop`

#3
Files changed (1) hide show
  1. 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[:, input_ids.shape[1] : input_ids.shape[1] + first_mask_position + 1]
 
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]