Add print statements

visual.py

@@ -106,6 +106,20 @@ class Transformer(nn.Module):
 
         for idx, layer_module in enumerate(self.layers):
             hidden_states = layer_module(hidden_states)
+
+        print("Shape of hidden states after CLIP:", hidden_states.shape)
+        torch.save(hidden_states, "hidden_states_after_clip.pt")
+
+        from huggingface_hub import HfApi
+
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj="hidden_states_after_clip.pt",
+            path_in_repo="hidden_states_after_clip.pt",
+            repo_id="nielsr/test-cogvlm",
+            repo_type="dataset",
+        )
+
         return hidden_states