Add print statements

modeling_cogvlm.py

@@ -241,10 +241,6 @@ 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:
-            print("Shape of query_states:", query_states.shape)
-            print("First values of query_states:", query_states[0,0,:3,:3])
-
         # if print_values:
 
             # torch.save(query_states, "query_states.pt")
@@ -287,6 +283,16 @@ class VisionExpertAttention(nn.Module):
 
         past_key_value = (key_states, value_states) if use_cache else None
 
+        if print_values:
+            print("Shape of query_states:", query_states.shape)
+            print("First values of query_states:", query_states[0,0,:3,:3])
+
+            print("Shape of key_states:", key_states.shape)
+            print("First values of key_states:", key_states[0,0,:3,:3])
+
+            print("Shape of value_states:", value_states.shape)
+            print("First values of value_states:", value_states[0,0,:3,:3])
+
         context_layer = attention_fn(
             query_layer=query_states, key_layer=key_states, value_layer=value_states, attention_mask=attention_mask,
             scaling_attention_score=True, attention_dropout=None)