diff --git a/README.md b/README.md
index 64e4ab6f597afcbf636c29884a8dcbab3c57642f..7b69a386f9a491b3806e9e6652607ea54fdf6fc4 100644
--- a/README.md
+++ b/README.md
@@ -141,9 +141,9 @@ In this work, we additionally use LOST predictions to train object detection mod
 The predictions of the class-agnostic Faster R-CNN model trained using LOST boxes as pseudo-gt are stored in the folder `data/CAD_predictions`. In order to launch the corloc evaluation, please launch the following scripts. It is to be noted that in this evaluation, only the box with the highest confidence score is considered per image. 
 
 ```
-python main_corloc_evaluation.py --dataset VOC07 --set trainval --type_predictions detectron --prediction_file data/CAD_predictions/LOST_plus_CAD_VOC07.json
-python main_corloc_evaluation.py --dataset VOC12 --set trainval --type_predictions detectron --prediction_file data/CAD_predictions/LOST_plus_CAD_VOC07.json
-python main_corloc_evaluation.py --dataset COCO --set train --type_predictions detectron --prediction_file data/CAD_predictions/LOST_plus_CAD_VOC07.json
+python main_corloc_evaluation.py --dataset VOC07 --set trainval --type_pred detectron --pred_file data/CAD_predictions/LOST_plus_CAD_VOC07.json
+python main_corloc_evaluation.py --dataset VOC12 --set trainval --type_pred detectron --pred_file data/CAD_predictions/LOST_plus_CAD_VOC12.json
+python main_corloc_evaluation.py --dataset COCO20k --set train --type_pred detectron --pred_file data/CAD_predictions/LOST_plus_CAD_COCO20k.json
 ```
 
 The following table presents the obtained corloc results.
diff --git a/main_corloc_evaluation.py b/main_corloc_evaluation.py
index c4262166a483f70ca613fc03dd93f859917d23d6..a4d9fb497b780d2d4fa1924db363b415732e7182 100755
--- a/main_corloc_evaluation.py
+++ b/main_corloc_evaluation.py
@@ -19,14 +19,14 @@ from datasets import Dataset, bbox_iou
 if __name__ == "__main__":
     parser = argparse.ArgumentParser("Visualize Self-Attention maps")
     parser.add_argument(
-        "--type_predictions",
+        "--type_pred",
         default="boxes_OD",
         choices=["boxes_OD", "detectron"],
         type=str,
         help="Type of predictions will inform on how to load",
     )
     parser.add_argument(
-        "--prediction_file", default="", type=str, help="File location of predictions"
+        "--pred_file", default="", type=str, help="File location of predictions."
     )
     parser.add_argument(
         "--dataset",
@@ -56,11 +56,14 @@ if __name__ == "__main__":
 
     # -------------------------------------------------------------------------------------------------------
     # Load predictions
-    if args.type_predictions == "boxes_OD":
-        with open(args.prediction_file, "rb") as f:
+    if not os.path.exists(args.pred_file):
+        raise ValueError(f"File {args.pred_file} does not exists.")
+
+    if args.type_pred == "boxes_OD":
+        with open(args.pred_file, "rb") as f:
             predictions = pickle.load(f)
-    elif args.type_predictions == "detectron":
-        with open(args.prediction_file, "r") as f:
+    elif args.type_pred == "detectron":
+        with open(args.pred_file, "r") as f:
             predictions = json.load(f)
 
     cnt = 0
@@ -87,9 +90,9 @@ if __name__ == "__main__":
             if gt_bbxs.shape[0] == 0 and args.no_hard:
                 continue
 
-        if args.type_predictions == "boxes_OD":
+        if args.type_pred == "boxes_OD":
             pred = np.asarray(predictions[im_name])
-        elif args.type_predictions == "detectron":
+        elif args.type_pred == "detectron":
             name_ind = im_name
             if "VOC" in args.dataset:
                 name_ind = im_name[:-4]