Add print statements

modeling_cogvlm.py

@@ -295,6 +295,8 @@ class CogVLMDecoderLayer(nn.Module):
 
         hidden_states = self.input_layernorm(hidden_states)
 
+        print("Hidden states before self attention:", hidden_states[0,:3,:3])
+
         # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
@@ -305,6 +307,9 @@ class CogVLMDecoderLayer(nn.Module):
             output_attentions=output_attentions,
             use_cache=use_cache,
         )
+
+        print("Hidden states after self attention:", hidden_states[0,:3,:3])
+
         hidden_states = residual + hidden_states
 
         # Fully Connected
@@ -522,7 +527,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
-            if idx in [0, 31]:
+            if idx in [0, 1, 2]:
                 print(f"Hidden states before layer {idx}", hidden_states[0,:3,:3])
 
             past_key_value = past_key_values[idx] if past_key_values is not None else None