diff --git a/src/data/dataloader.py b/src/data/dataloader.py
index c1f444a..c6adff8 100644
--- a/src/data/dataloader.py
+++ b/src/data/dataloader.py
@@ -46,7 +46,8 @@ class Spheres(pl.LightningDataModule):
             ),
         )
 
-        dataset = LabeledDataset(image_dir="/dev/shm/TRAIN/", transforms=transforms)
+        # dataset = LabeledDataset(image_dir="/dev/shm/TRAIN/", transforms=transforms)
+        dataset = LabeledDataset(image_dir=wandb.config.DIR_TRAIN_IMG, transforms=transforms)
         # dataset = Subset(dataset, range(6 * 200))  # subset for debugging purpose
         # dataset = Subset(dataset, [0] * 320)  # overfit test
 
@@ -84,7 +85,8 @@ class Spheres(pl.LightningDataModule):
             ),
         )
 
-        dataset = RealDataset(root="/dev/shm/TEST/", transforms=transforms)
+        # dataset = RealDataset(root="/dev/shm/TEST/", transforms=transforms)
+        dataset = RealDataset(root=wandb.config.DIR_VALID_IMG, transforms=transforms)
 
         return DataLoader(
             dataset,
diff --git a/wandb.yaml b/wandb.yaml
index da6ca75..44ea6b9 100644
--- a/wandb.yaml
+++ b/wandb.yaml
@@ -1,7 +1,7 @@
-# DIR_TRAIN_IMG:
-#   value: "/media/disk1/lfainsin/BACKGROUND/"
-# DIR_VALID_IMG:
-#   value: "/media/disk1/lfainsin/TEST_batched/"
+DIR_TRAIN_IMG:
+  value: "/media/disk1/lfainsin/TRAIN_prerender/"
+DIR_VALID_IMG:
+  value: "/media/disk1/lfainsin/TEST_tmp_mrcnn/"
 # DIR_SPHERE:
 #   value: "/media/disk1/lfainsin/SPHERES/"