Add print statements

modeling_cogvlm.py

@@ -290,6 +290,11 @@ class VisionExpertAttention(nn.Module):
         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)
+        
+        if print_values:
+            print("Shape of context_layer:", context_layer.shape)
+            print("First values of context_layer:", context_layer[0,0,:3,:3])
+        
         if context_layer.size() != (bsz, self.num_heads, q_len, self.head_dim):
             raise ValueError(
                 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
@@ -657,7 +662,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
                 past_key_value=past_key_value,
                 output_attentions=output_attentions,
                 use_cache=use_cache,
-                print_values=idx==0,
+                print_values=idx==0 and step==1,
             )
             hidden_states = layer_outputs[0]