diff --git a/main.py b/main.py
index 4047c2db426db3a70c255a9bcaa3268829a8ffb1..9ccc58d94033046f2911d91937de452af12e2d64 100644
--- a/main.py
+++ b/main.py
@@ -59,6 +59,7 @@ if __name__ == '__main__':
     # Add an argument
     parser.add_argument('--Nb_worlds', type=int, required=True)
     parser.add_argument('--World_begin', type=int, required=True)
+    parser.add_argument('--dataset_id', type=str, required=True)
     # Parse the argument
     args = parser.parse_args()
 
@@ -69,7 +70,7 @@ if __name__ == '__main__':
     Nb_instance = 1
     occ_target = 0.5
 
-    dataset_src = "/gpfsscratch/rech/uli/ubn15wo/dataset/s2rg/Fruits_all_medium/data/"
+    dataset_src = f"/gpfsscratch/rech/uli/ubn15wo/dataset{args.dataset_id}/s2rg/Fruits_all_medium/data/"
     #dataset_src = "/media/mahmoud/E/Fruits_easy/data"
     #dataset_src = "/media/gduret/DATA/dataset/s2rg/Fruits_all_medium/data"
 
@@ -77,7 +78,7 @@ if __name__ == '__main__':
     data_options = {"high": "ground_truth_rgb",
                     "low": "ground_truth_depth"}
     dataset_type = data_options[choice]
-    dataset_name = f"/gpfsscratch/rech/uli/ubn15wo/dataset/s2rg/Fruits_all_medium/GUIMOD_{choice}"
+    dataset_name = f"/gpfsscratch/rech/uli/ubn15wo/dataset{args.dataset_id}/s2rg/Fruits_all_medium/GUIMOD_{choice}"
     list_categories = ["banana1", "kiwi1", "pear2", "apricot", "orange2", "peach1", "lemon2", "apple2" ]
     Nb_camera = 15
     #Nb_world = 10000