change torch.IntTensor to torch.LongTensor

lufficc · May 22, 2019 · c5b0069 · c5b0069
1 parent b14cad4
commit c5b0069
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ssd/utils/distributed_util.py b/ssd/utils/distributed_util.py
@@ -71,8 +71,8 @@ def all_gather(data):
     tensor = torch.ByteTensor(storage).to("cuda")
 
     # obtain Tensor size of each rank
-    local_size = torch.IntTensor([tensor.numel()]).to("cuda")
-    size_list = [torch.IntTensor([0]).to("cuda") for _ in range(world_size)]
+    local_size = torch.LongTensor([tensor.numel()]).to("cuda")
+    size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)]
     dist.all_gather(size_list, local_size)
     size_list = [int(size.item()) for size in size_list]
     max_size = max(size_list)