This commit is contained in:
Glenn Jocher
2020-03-02 14:28:08 -08:00
parent 44daace4ca
commit 2774c1b398
+2 -1
View File
@@ -55,7 +55,8 @@ def test(cfg,
# Dataloader # Dataloader
if dataloader is None: if dataloader is None:
dataset = LoadImagesAndLabels(path, img_size, batch_size, rect=True) dataset = LoadImagesAndLabels(path, img_size, batch_size, rect=True, single_cls=opt.single_cls,
cache_labels=True)
batch_size = min(batch_size, len(dataset)) batch_size = min(batch_size, len(dataset))
dataloader = DataLoader(dataset, dataloader = DataLoader(dataset,
batch_size=batch_size, batch_size=batch_size,