Merge pull request #157 from annabelle1217/annabelle/fix-pytorch-unet…

…-load ✏️ typo fix for pycaret component library
XpressAI · May 11, 2022 · 8ad45a5 · 8ad45a5
2 parents 282a360 + 590ee6a
commit 8ad45a5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xai_components/xai_pytorch/training.py b/xai_components/xai_pytorch/training.py
@@ -237,7 +237,7 @@ def execute(self, ctx) -> None:
                                                    pin_memory=pin_memory_bool
                                                    )
 
-        tests_set = UNetDataset(image_dir=str(self.train_image_folder.value[0]), masks_dir=str(self.train_image_folder.value[1]))
+        tests_set = UNetDataset(image_dir=str(self.test_image_folder.value[0]), masks_dir=str(self.test_image_folder.value[1]))
         tests_loader = torch.utils.data.DataLoader(tests_set,
                                                    batch_size=image_batch_size,
                                                    shuffle=shuffle_bool,