update modeling_yi.py

modeling_yi.py

@@ -539,7 +539,7 @@ class YiModel(YiPreTrainedModel):
     def _prepare_decoder_attention_mask(
         self, attention_mask, input_ids, inputs_embeds, past_key_values_length
     ):
-        input_shape = input_ids.shape
+        input_shape = input_ids.shape if input_ids else inputs_embeds.shape[:-1]
         # create causal mask
         # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
         combined_attention_mask = None