diff --git a/modeling_internvl_chat.py b/modeling_internvl_chat.py
index cb83c57..ccf2ab9 100644
--- a/modeling_internvl_chat.py
+++ b/modeling_internvl_chat.py
@@ -123,8 +123,8 @@ class InternVLChatModel(PreTrainedModel):
             vit_embeds = vit_embeds.reshape(-1, C)
             print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
                   f'vit_embeds.shape={vit_embeds.shape}')
-            n_token = selected.sum()
-            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]
+            n_token = min(selected.sum(), vit_embeds.size(0))
+            input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]
 
         input_embeds = input_embeds.reshape(B, N, C)