Update modeling_internlm2_ve.py

modeling_internlm2_ve.py

@@ -689,20 +689,18 @@ class InternLM2DecoderLayer(nn.Module):
         hidden_states = self.ffn_norm(hidden_states)
 
         if past_key_value is None:
-            """
-            *************************
-            maybe faster
-            ***************************
-            """
-            ##############################################################################################################
-            # dim=hidden_states.shape[-1]
-            # visual_token_mask=visual_token_mask.repeat(1,1,dim)
-            # if visual_token_mask.any():
-            #     hidden_states[visual_token_mask] = self.feed_forward_ve(hidden_states[visual_token_mask].reshape(-1,dim)).reshape(-1)
-            # if (~visual_token_mask).any(): 
-            #     hidden_states[~visual_token_mask] = self.feed_forward(hidden_states[~visual_token_mask].reshape(-1,dim)).reshape(-1)
+            ##########################################--modified by luogen--##############################################
+            if self.training:
+                hidden_states = self.feed_forward(hidden_states)*(1.-visual_token_mask)+ self.feed_forward_ve(hidden_states)*visual_token_mask
+            else:
+                dim=hidden_states.shape[-1]
+                visual_token_mask=visual_token_mask.repeat(1,1,dim).bool()
+                non_visual_token_mask=~visual_token_mask
+                if visual_token_mask.any():
+                    hidden_states[visual_token_mask] = self.feed_forward_ve(hidden_states[visual_token_mask].reshape(-1,dim)).reshape(-1)
+                if (non_visual_token_mask).any(): 
+                    hidden_states[non_visual_token_mask] = self.feed_forward(hidden_states[non_visual_token_mask].reshape(-1,dim)).reshape(-1)
             ##############################################################################################################
-            hidden_states = self.feed_forward(hidden_states)*(1.-visual_token_mask)+ self.feed_forward_ve(hidden_states)*visual_token_mask
         else:
             hidden_states = self.feed_forward(hidden_states)