updates

2020-03-02 14:28:08 -08:00
parent 44daace4ca
commit 2774c1b398
1 changed files with 2 additions and 1 deletions
@@ -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,