Add print statements

modeling_cogvlm.py

@@ -522,7 +522,8 @@ class CogVLMModel(CogVLMPreTrainedModel):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
-            print(f"Hidden states before layer {idx}", hidden_states[0,:3,:3])
+            if idx in [0, 31]:
+                print(f"Hidden states before layer {idx}", hidden_states[0,:3,:3])
 
             past_key_value = past_key_values[idx] if past_key_values is not None else None
             layer_outputs = decoder_layer(