Add print statements

visual.py

@@ -90,7 +90,21 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([TransformerLayer(config) for _ in range(config.num_hidden_layers)])
 
     def forward(self, hidden_states):
-        for layer_module in self.layers:
+
+        print("Shape of hidden states before CLIP:", hidden_states.shape)
+        torch.save(hidden_states, "hidden_states_before_clip.pt")
+
+        from huggingface_hub import HfApi
+
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj="hidden_states_before_clip.pt",
+            path_in_repo="hidden_states_before_clip.pt",
+            repo_id="nielsr/test-cogvlm",
+            repo_type="dataset",
+        )
+
+        for idx, layer_module in enumerate(self.layers):
             hidden_states = layer_module(hidden_states)
         return hidden_states