Add print statements

modeling_cogvlm.py

@@ -241,33 +241,35 @@ class VisionExpertAttention(nn.Module):
         key_states = self._transpose_for_scores(key_states)  # B, H, L, HD
         value_states = self._transpose_for_scores(value_states)  # B, H, L, HD
 
-        torch.save(query_states, "query_states.pt")
-        torch.save(key_states, "key_states.pt")
-        torch.save(value_states, "value_states.pt")
+        if print_values:
 
-        from huggingface_hub import HfApi
+            torch.save(query_states, "query_states.pt")
+            torch.save(key_states, "key_states.pt")
+            torch.save(value_states, "value_states.pt")
 
-        api = HfApi()
-        api.upload_file(
-            path_or_fileobj="query_states.pt",
-            path_in_repo="query_states.pt",
-            repo_id="nielsr/test-cogvlm",
-            repo_type="dataset",
-        )
-        api = HfApi()
-        api.upload_file(
-            path_or_fileobj="key_states.pt",
-            path_in_repo="key_states.pt",
-            repo_id="nielsr/test-cogvlm",
-            repo_type="dataset",
-        )
-        api = HfApi()
-        api.upload_file(
-            path_or_fileobj="value_states.pt",
-            path_in_repo="value_states.pt",
-            repo_id="nielsr/test-cogvlm",
-            repo_type="dataset",
-        )
+            from huggingface_hub import HfApi
+
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj="query_states.pt",
+                path_in_repo="query_states.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj="key_states.pt",
+                path_in_repo="key_states.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj="value_states.pt",
+                path_in_repo="value_states.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
 
         kv_seq_len = key_states.shape[-2]
         if past_key_value is not None: