Add print statements

modeling_cogvlm.py

@@ -296,8 +296,8 @@ class CogVLMDecoderLayer(nn.Module):
 
         hidden_states = self.input_layernorm(hidden_states)
 
-        # if print_values:
-        #     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(
@@ -310,8 +310,8 @@ class CogVLMDecoderLayer(nn.Module):
             use_cache=use_cache,
         )
 
-        # if print_values:
-        #     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
 
@@ -600,10 +600,6 @@ class CogVLMModel(CogVLMPreTrainedModel):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
-            # if idx in [0, 1, 2]:
-            #     print(f"Hidden states before layer {idx}", hidden_states[0,:3,:3])
-            #     print(f"Mean of hidden states before layer {idx}", hidden_states.mean())
-
             past_key_value = past_key_values[idx] if past_key_values is not None else None
             layer_outputs = decoder_layer(
                 hidden_states,
@@ -613,7 +609,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],
+                print_values=idx==0,
             )
             hidden_states = layer_outputs[0]