Add print statements

modeling_cogvlm.py

@@ -241,7 +241,7 @@ class VisionExpertAttention(nn.Module):
         key_states = self._transpose_for_scores(key_states)  # B, H, L, HD
         value_states = self._transpose_for_scores(value_states)  # B, H, L, HD
 
-        if print_values:
+        # if print_values:
 
             # torch.save(query_states, "query_states.pt")
             # torch.save(key_states, "key_states.pt")
@@ -325,13 +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])
+        # 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 after RMS norm, before self attention:", hidden_states[0,:3,:3])
+        # if print_values:
+        #     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(
@@ -345,8 +345,8 @@ class CogVLMDecoderLayer(nn.Module):
             print_values=print_values,
         )
 
-        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