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

* fix: quad training

* fix: quad training in segmentation
pull/13365/head
imyhxy 2024-10-16 00:22:02 +08:00 committed by GitHub
parent abdfbd68d1
commit 94a6245617
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 0 deletions

View File

@ -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,

View File

@ -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,