diff --git a/image_ref/main_sweep.py b/image_ref/main_sweep.py
index a9e169659ba38a924a6c69f4a2bfd0193a934566..64fb7c7a550b39ebc693e0f512c14871c6436354 100644
--- a/image_ref/main_sweep.py
+++ b/image_ref/main_sweep.py
@@ -1,10 +1,8 @@
+import time
 
 import wandb as wdb
-from config import load_args_contrastive
-
 
 if __name__ == '__main__':
-    args = load_args_contrastive()
     sweep_configuration = {
         "program": "sweep_train.py",
         "method": "random",
@@ -23,9 +21,15 @@ if __name__ == '__main__':
             "dataset_val_dir": {"value": "data/processed_data_wiff/npy_image/test_data"},
             "dataset_ref_dir": {"values": ["image_ref/img_ref","image_ref/img_ref_count_th_10","image_ref/img_ref_count_th_5"]},
         },
+        "controller":{
+            "type": "local"},
     }
-    sweep_id = wdb.sweep(sweep=sweep_configuration, project="param_sweep_contrastive")
 
+    sweep_id = wdb.sweep(sweep_configuration)
+
+    # Start the local controller
     sweep = wdb.controller(sweep_id)
-    sweep.configure_controller(type="local")
-    sweep.run(verbose=True,print_debug=True,print_status=True,print_actions=True)
+    while not sweep.done():
+        sweep.print_status()
+        sweep.step()
+        time.sleep(5)