Update distributed.py

microsoft · Oct 12, 2024 · a2efda0 · a2efda0
1 parent 270bf31
commit a2efda0
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/getmusic/getmusic/distributed/distributed.py b/getmusic/getmusic/distributed/distributed.py
@@ -81,14 +81,14 @@ def all_gather(data):
     if world_size == 1:
         return [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 = None 
     # buffer = pickle.dumps(data) 
+
+    try:
+        storage = torch.ByteStorage.from_buffer(buffer)
+    except Exception as e:
+        raise RuntimeError("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 line 85 of getmusic/distributed/distributed.py")
 
-    storage = torch.ByteStorage.from_buffer(buffer)
     tensor = torch.ByteTensor(storage).to("cuda")
 
     local_size = torch.IntTensor([tensor.numel()]).to("cuda")