Skip to content

Commit

Permalink
Merge pull request #4 from NanoCode012/patch-4
Browse files Browse the repository at this point in the history
Add drop_last for multi gpu
  • Loading branch information
MagicFrogSJTU authored Jul 8, 2020
2 parents b9a50ae + 02c63ef commit 4b8450b
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions utils/datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=Fa
batch_size=batch_size,
num_workers=nw,
sampler=train_sampler,
drop_last=(train_sampler is not None),
pin_memory=True,
collate_fn=LoadImagesAndLabels.collate_fn)
return dataloader, dataset
Expand Down

0 comments on commit 4b8450b

Please sign in to comment.