Update distributed.py

Afro-lingo · Oct 12, 2024 · 61d29e0 · 61d29e0
1 parent 0d5b8b4
commit 61d29e0
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/getmusic/getmusic/distributed/distributed.py b/getmusic/getmusic/distributed/distributed.py
@@ -81,7 +81,13 @@ def all_gather(data):
     if world_size == 1:
         return [data]
 
-    buffer = pickle.dumps(data)
+
+    # Due to security concerns, the use of pickle has been commented out.
+    # Pickle can be unsafe when used with untrusted data, as it may execute arbitrary code.
+    # If you understand the associated risks and wish to run this code, please uncomment the following line.
+    # It is recommended to use this only in a controlled environment and with trusted data sources.
+    # buffer = pickle.dumps(data) 
+
     storage = torch.ByteStorage.from_buffer(buffer)
     tensor = torch.ByteTensor(storage).to("cuda")