diff --git a/comp.ipynb.REMOVED.git-id b/comp.ipynb.REMOVED.git-id
index 3c6779c..6ca72da 100644
--- a/comp.ipynb.REMOVED.git-id
+++ b/comp.ipynb.REMOVED.git-id
@@ -1 +1 @@
-3c9a34f197340a6051eb34d11695c7d6b72164f0
\ No newline at end of file
+5ef2ef54312186cd3e3162869c4f237b69de3b1e
\ No newline at end of file
diff --git a/src/train.py b/src/train.py
index eb1b115..1f72d51 100644
--- a/src/train.py
+++ b/src/train.py
@@ -13,13 +13,12 @@ CONFIG = {
     "DIR_TRAIN_IMG": "/home/lilian/data_disk/lfainsin/train/",
     "DIR_VALID_IMG": "/home/lilian/data_disk/lfainsin/val/",
     "DIR_TEST_IMG": "/home/lilian/data_disk/lfainsin/test/",
-    "DIR_SPHERE_IMG": "/home/lilian/data_disk/lfainsin/spheres/Images/",
-    "DIR_SPHERE_MASK": "/home/lilian/data_disk/lfainsin/spheres/Masks/",
+    "DIR_SPHERE": "/home/lilian/data_disk/lfainsin/spheres_prod/",
     # "FEATURES": [1, 2, 4, 8],
     # "FEATURES": [4, 8, 16, 32],
-    # "FEATURES": [8, 16, 32, 64],
+    "FEATURES": [8, 16, 32, 64],
     # "FEATURES": [4, 8, 16, 32, 64],
-    "FEATURES": [8, 16, 32, 64, 128],
+    # "FEATURES": [8, 16, 32, 64, 128],
     # "FEATURES": [16, 32, 64, 128],
     # "FEATURES": [64, 128, 256, 512],
     "N_CHANNELS": 3,
diff --git a/src/unet/model.py b/src/unet/model.py
index 4508b27..8d5288f 100644
--- a/src/unet/model.py
+++ b/src/unet/model.py
@@ -73,7 +73,7 @@ class UNet(pl.LightningModule):
                 A.Resize(wandb.config.IMG_SIZE, wandb.config.IMG_SIZE),
                 A.Flip(),
                 A.ColorJitter(),
-                RandomPaste(wandb.config.SPHERES, wandb.config.DIR_SPHERE_IMG, wandb.config.DIR_SPHERE_MASK),
+                RandomPaste(wandb.config.SPHERES, wandb.config.DIR_SPHERE),
                 A.GaussianBlur(),
                 A.ISONoise(),
                 A.ToFloat(max_value=255),
diff --git a/src/utils/paste.py b/src/utils/paste.py
index a25289b..0c4db00 100644
--- a/src/utils/paste.py
+++ b/src/utils/paste.py
@@ -1,4 +1,3 @@
-import os
 import random as rd
 from pathlib import Path
 
@@ -75,20 +74,10 @@ class RandomPaste(A.DualTransform):
         mask_path = img_path.parent.joinpath("MASK.PNG")
 
         # load the "paste" image
-        paste_img = Image.open(
-            os.path.join(
-                self.path_paste_img_dir,
-                img_path,
-            )
-        ).convert("RGBA")
+        paste_img = Image.open(img_path).convert("RGBA")
 
         # load its respective mask
-        paste_mask = Image.open(
-            os.path.join(
-                self.path_paste_mask_dir,
-                mask_path,
-            )
-        ).convert("LA")
+        paste_mask = Image.open(mask_path).convert("LA")
 
         # load the target image
         target_img = params["image"]
@@ -151,6 +140,3 @@ class RandomPaste(A.DualTransform):
         )
 
         return params
-
-    def get_transform_init_args_names(self):
-        return "scale_range", "path_paste_img_dir", "path_paste_mask_dir"