From 94a6245617adf3df848e4879458f29208c50fa1f Mon Sep 17 00:00:00 2001
From: imyhxy <imyhxy@gmail.com>
Date: Wed, 16 Oct 2024 00:22:02 +0800
Subject: [PATCH] Fix bug when enable `--quad` training option (#13355)
* fix: quad training
* fix: quad training in segmentation
---
utils/dataloaders.py | 1 +
utils/segment/dataloaders.py | 1 +
2 files changed, 2 insertions(+)
diff --git a/utils/dataloaders.py b/utils/dataloaders.py
index 61358eb9f..00e98168f 100644
--- 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
index 5f5666c38..f65076494 100644
--- 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,