Merge branch 'feature/DDP_fixed' of https://github.com/MagicFrogSJTU/…

…yolov5 into feature/DDP_fixed
fhkiel-mlaip · Jul 9, 2020 · 1dabe33 · 1dabe33
2 parents a1ce9b1 + 4b8450b
commit 1dabe33
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/utils/datasets.py b/utils/datasets.py
@@ -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