Add print statements

modeling_cogvlm.py

@@ -325,10 +325,13 @@ class CogVLMDecoderLayer(nn.Module):
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         residual = hidden_states
 
+        if print_values:
+            print("Hidden states before RMS norm:", hidden_states[0, :3, :3])
+
         hidden_states = self.input_layernorm(hidden_states)
 
         if print_values:
-            print("Hidden states before self attention:", hidden_states[0,:3,:3])
+            print("Hidden states after RMS norm, before self attention:", hidden_states[0,:3,:3])
 
         # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(