diff --git a/sort.py b/sort.py
new file mode 100644
index 0000000..e69de29
diff --git a/train.py b/train.py
index 5395af7..4856f1f 100644
--- a/train.py
+++ b/train.py
@@ -78,7 +78,7 @@ class GarageDoorCNN(nn.Module):
 
 def train_model():
     # --- Hyperparameters and Configuration ---
-    DATA_DIR = 'data'
+    DATA_DIR = 'data/labelled'
     MODEL_SAVE_PATH = 'garage_door_cnn.pth'
     NUM_EPOCHS = 10
     BATCH_SIZE = 32
@@ -158,7 +158,7 @@ def train_model():
 
 if __name__ == '__main__':
     # Check if data directory exists
-    if not os.path.isdir('data/open') or not os.path.isdir('data/closed'):
+    if not os.path.isdir('data/labelled/open') or not os.path.isdir('data/labelled/closed'):
         print("Error: Data directories 'data/open' and 'data/closed' not found.")
         print("Please create them and place your image snapshots inside.")
     else: