diff --git a/src/config.py b/src/config.py
index 1e0eaa2..8651565 100644
--- a/src/config.py
+++ b/src/config.py
@@ -1,3 +1,4 @@
+import random
 from pathlib import Path
 
 DATA = sorted(Path("/mnt/wsl/PHYSICALDRIVE1/data/unsplash").glob("*.jpg"))
@@ -9,10 +10,18 @@ LOGS_PATH = Path("/home/andras/projects/bipolaroid/logs")
 RUNS_PATH = Path("/home/andras/projects/bipolaroid/runs")
 
 
-for path in [
-    CACHE_PATH,
-    MODELS_PATH,
-    LOGS_PATH,
-    RUNS_PATH
-]:
+for path in [CACHE_PATH, MODELS_PATH, LOGS_PATH, RUNS_PATH]:
     path.mkdir(exist_ok=True, parents=True)
+
+
+length = len(DATA)
+indices = list(range(length))
+
+random.seed(42)
+random.shuffle(indices)
+
+train_indices = indices[: int(length * TRAIN_SIZE)]
+test_indices = indices[int(length * TRAIN_SIZE) :]
+
+TRAIN_DATA = [DATA[i] for i in train_indices]
+TEST_DATA = [DATA[i] for i in test_indices]
diff --git a/src/data.py b/src/data.py
deleted file mode 100644
index f29406b..0000000
--- a/src/data.py
+++ /dev/null
@@ -1,14 +0,0 @@
-import random
-from config import DATA, TRAIN_SIZE
-
-random.seed(42)
-length = len(DATA)
-indices = list(range(length))
-
-random.shuffle(indices)
-
-train_indices = indices[: int(length * TRAIN_SIZE)]
-test_indices = indices[int(length * TRAIN_SIZE) :]
-
-TRAIN_DATA = [DATA[i] for i in train_indices]
-TEST_DATA = [DATA[i] for i in test_indices]