Fix bug when enable --quad training option (#13355)

* fix: quad training * fix: quad training in segmentation
ultralytics · Oct 15, 2024 · 94a6245 · 94a6245
1 parent abdfbd6
commit 94a6245
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/utils/dataloaders.py b/utils/dataloaders.py
@@ -210,6 +210,7 @@ def create_dataloader(
         shuffle=shuffle and sampler is None,
         num_workers=nw,
         sampler=sampler,
+        drop_last=quad,
         pin_memory=PIN_MEMORY,
         collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn,
         worker_init_fn=seed_worker,

diff --git a/utils/segment/dataloaders.py b/utils/segment/dataloaders.py
@@ -75,6 +75,7 @@ def create_dataloader(
         shuffle=shuffle and sampler is None,
         num_workers=nw,
         sampler=sampler,
+        drop_last=quad,
         pin_memory=True,
         collate_fn=LoadImagesAndLabelsAndMasks.collate_fn4 if quad else LoadImagesAndLabelsAndMasks.collate_fn,
         worker_init_fn=seed_worker,