Fix Views (#47)

luxonis · Jul 17, 2024 · ef606b7 · ef606b7
1 parent 57258ee
commit ef606b7
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py
@@ -128,11 +128,11 @@ def __init__(
                     if view == "train"
                     else self.val_augmentations
                 ),
-                view=(
-                    self.cfg.loader.train_view
-                    if view == "train"
-                    else self.cfg.loader.val_view
-                ),
+                view={
+                    "train": self.cfg.loader.train_view,
+                    "val": self.cfg.loader.val_view,
+                    "test": self.cfg.loader.test_view,
+                }[view],
                 image_source=self.cfg.loader.image_source,
                 **self.cfg.loader.params,
             )