Add print statements

modeling_cogvlm.py

@@ -290,12 +290,14 @@ class CogVLMDecoderLayer(nn.Module):
             past_key_value: Optional[Tuple[torch.Tensor]] = None,
             output_attentions: Optional[bool] = False,
             use_cache: Optional[bool] = False,
+            print_values = False,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         residual = hidden_states
 
         hidden_states = self.input_layernorm(hidden_states)
 
-        print("Hidden states before self attention:", hidden_states[0,:3,:3])
+        if print_values:
+            print("Hidden states before self attention:", hidden_states[0,:3,:3])
 
         # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
@@ -308,7 +310,8 @@ class CogVLMDecoderLayer(nn.Module):
             use_cache=use_cache,
         )
 
-        print("Hidden states after self attention:", hidden_states[0,:3,:3])
+        if print_values:
+            print("Hidden states after self attention:", hidden_states[0,:3,:3])
 
         hidden_states = residual + hidden_states
 
@@ -539,6 +542,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
                 past_key_value=past_key_value,
                 output_attentions=output_attentions,
                 use_cache=use_cache,
+                print_values=idx in [0, 1, 2],
             )
             hidden_states = layer_outputs[0]