Fix warning

BloodAxe · May 4, 2023 · bf29472 · bf29472
1 parent 45fafad
commit bf29472
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_toolbelt/utils/distributed.py b/pytorch_toolbelt/utils/distributed.py
@@ -158,7 +158,7 @@ def all_gather(data: Any) -> List[Any]:
 
     # serialized to a Tensor
     buffer = pickle.dumps(data)
-    storage = torch.ByteStorage.from_buffer(buffer)
+    storage = torch.UntypedStorage.from_buffer(buffer, dtype=torch.uint8)
     tensor = torch.ByteTensor(storage).to("cuda")
 
     # obtain Tensor size of each rank