Fix deprecation warning of using TypedStorage

BloodAxe · Aug 21, 2023 · 30e3a05 · 30e3a05
1 parent 8427960
commit 30e3a05
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_toolbelt/utils/distributed.py b/pytorch_toolbelt/utils/distributed.py
@@ -137,7 +137,7 @@ def broadcast_from_master(data: Any) -> Any:
 
     if local_rank == 0:
         buffer = pickle.dumps(data)
-        storage = torch.ByteStorage.from_buffer(buffer)
+        storage = torch.UntypedStorage.from_buffer(buffer, dtype=torch.uint8)
         payload = torch.ByteTensor(storage).to("cuda")
         local_size = payload.numel()
     else: