zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

torch.from_tensor() bug fix

This commit is contained in:
Glenn Jocher 2020-08-25 03:14:17 -07:00
parent 4fb8cb353f
commit 09402a2174
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
train.py
View File

@ -225,7 +225,7 @@ def train(hyp, opt, device, tb_writer=None):
if rank != -1:
indices = torch.zeros([dataset.n], dtype=torch.int)
if rank == 0:
indices[:] = torch.from_tensor(dataset.indices, dtype=torch.int)
indices[:] = torch.tensor(dataset.indices, dtype=torch.int)
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()