Add print statements

modeling_cogvlm.py

@@ -285,13 +285,16 @@ class VisionExpertAttention(nn.Module):
 
         if print_values:
             print("Shape of query_states:", query_states.shape)
-            print("First values of query_states:", query_states[0,0,:3,:3])
+            print("Last values of query_states:", query_states[0,0,-3:,-3:])
+            print("Mean of query_states:", query_states.mean())
 
             print("Shape of key_states:", key_states.shape)
-            print("First values of key_states:", key_states[0,0,:3,:3])
+            print("Last values of key_states:", key_states[0,0,-3:,-3:])
+            print("Mean of key_states:", key_states.mean())
 
             print("Shape of value_states:", value_states.shape)
-            print("First values of value_states:", value_states[0,0,:3,:3])
+            print("First values of value_states:", value_states[0,0,-3:,-3:])
+            print("Mean of value_states:", value_states.mean())
 
         context_layer = attention_fn(
             query_layer=query_states, key_layer=key_states, value_layer=value_states, attention_mask=attention_mask,