diff --git a/image_ref/main_sweep.py b/image_ref/main_sweep.py
index f7c47301881348d0fc1fba00cacc5df9dcdd4eb5..ac6d761775055da458d7848220039b5e7b1ac4ec 100644
--- a/image_ref/main_sweep.py
+++ b/image_ref/main_sweep.py
@@ -21,7 +21,7 @@ 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"]},
         },
-        "max_iter": 10
+        "max_iter": 100
     }
 
     for i in range(sweep_configuration["max_iter"]):
@@ -37,5 +37,6 @@ if __name__ == '__main__':
                     run_config[p] = random.uniform(v["min"],v["max"])
                 elif v["distribution"]=="log_uniform_values":
                     run_config[p] = np.exp(random.uniform(np.log(v["min"]), np.log(v["max"])))
-        print('Launching run')
+        print('Launching run ',i)
+        print('Parameter : ', sweep_configuration["parameters"])
         run_duo(run_config)
diff --git a/image_ref/sweep_train.py b/image_ref/sweep_train.py
index d2a273ae07cbd75bde1624e00d89303a709c751b..418256a6d1339f222748e7496c6fc922c74bfb04 100644
--- a/image_ref/sweep_train.py
+++ b/image_ref/sweep_train.py
@@ -90,7 +90,7 @@ def run_duo(args):
     os.environ["WANDB_MODE"] = "offline"
     os.environ["WANDB_DIR"] = os.path.abspath("./wandb_run")
 
-    wdb.init(project="param_sweep_contrastive", dir='./wandb_run')
+    wdb.init(project="param_sweep_contrastive", dir='./wandb_run', config=args)
 
     print('Wandb initialised')
     # load data