diff --git a/test.py b/test.py index 53fa7836..caa4284c 100644 --- a/test.py +++ b/test.py @@ -55,7 +55,8 @@ def test(cfg, # Dataloader 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)) dataloader = DataLoader(dataset, batch_size=batch_size,