Add print statements

modeling_cogvlm.py

@@ -676,6 +676,10 @@ class CogVLMModel(CogVLMPreTrainedModel):
 
         hidden_states = self.norm(hidden_states)
 
+        if step == 1:
+            print("Shape of hidden states:", hidden_states.shape)
+            print("First values of hidden states:", hidden_states[0,:3,:3])
+
         # add hidden states from the last decoder layer
         if output_hidden_states:
             all_hidden_states += (hidden_states,)