Skip to content
This repository was archived by the owner on Oct 31, 2023. It is now read-only.

Commit 16b8b09

Browse files
caiqifmassa
authored andcommitted
change torch.IntTensor to torch.LongTensor for storing large tensor numel (#799)
1 parent 1127bdd commit 16b8b09

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

maskrcnn_benchmark/utils/comm.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,8 @@ def all_gather(data):
6363
tensor = torch.ByteTensor(storage).to("cuda")
6464

6565
# obtain Tensor size of each rank
66-
local_size = torch.IntTensor([tensor.numel()]).to("cuda")
67-
size_list = [torch.IntTensor([0]).to("cuda") for _ in range(world_size)]
66+
local_size = torch.LongTensor([tensor.numel()]).to("cuda")
67+
size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)]
6868
dist.all_gather(size_list, local_size)
6969
size_list = [int(size.item()) for size in size_list]
7070
max_size = max(size_list)

0 commit comments

Comments
 (0)