diff --git a/README.md b/README.md
index 66cbd011c1dad1a52195b9ff8aa1f85ae2d37af3..0fb1abc7bf9df4bd5ccc77bcaaa136bb49070605 100644
--- a/README.md
+++ b/README.md
@@ -8,6 +8,8 @@ touch __init__.py
 echo -e "import sys\nfrom os.path import dirname, join\nsys.path.insert(0, join(dirname(__file__), '.'))" >> __init__.py;
 cd -
 
+# make sure coco dataset is present at $LOST_PATH/datasets/COCO
+ln -sf /root/coco /root/lost/datasets/COCO
 # replace $COCO_ROOT with path to COCO
 # replace * with the script you want to run
 COCO_ROOT=$COCO_ROOT ./scripts/run_*.sh
diff --git a/datasets/COCO b/datasets/COCO
new file mode 120000
index 0000000000000000000000000000000000000000..6385f3b3295789558a568e1966d0562b657bb151
--- /dev/null
+++ b/datasets/COCO
@@ -0,0 +1 @@
+/root/coco
\ No newline at end of file
diff --git a/main_lost.py b/main_lost.py
index f1def0e3495e7214630491a7277be8cd74473238..49f8cb239dafc4c689949b887a5f687579b4725c 100755
--- a/main_lost.py
+++ b/main_lost.py
@@ -133,6 +133,7 @@ if __name__ == "__main__":
     # -------------------------------------------------------------------------------------------------------
     # Model
     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    print("Running on device:", device)
     model = get_model(args.arch, args.patch_size, args.resnet_dilate, device)
 
     # -------------------------------------------------------------------------------------------------------
@@ -161,7 +162,7 @@ if __name__ == "__main__":
     if args.visualize:
         vis_folder = f"{args.output_dir}/visualizations/{exp_name}"
         os.makedirs(vis_folder, exist_ok=True)
-
+    
     # -------------------------------------------------------------------------------------------------------
     # Loop over images
     preds_dict = {}
diff --git a/scripts/run-dataset.sh b/scripts/run-dataset.sh
index fa2a485210d9db7118f499feac02793471a40c84..1a105092370150274ebf8d8a12d899de5d6e2a6f 100755
--- a/scripts/run-dataset.sh
+++ b/scripts/run-dataset.sh
@@ -1,5 +1,5 @@
 
-OUTPUT_PATH=/root/lost/outputs/kitti
+OUTPUT_PATH=/root/lost/outputs/coco
 
 DINO_ARCH=vit_base
 LOST_FEATURES=k
@@ -12,7 +12,7 @@ mkdir -p $OUTPUT_PATH
 echo $OUTPUT_PATH
 
 python main_lost.py \
-    --dataset KITTI \
+    --dataset COCO20k \
     --output_dir $OUTPUT_PATH \
     --arch $DINO_ARCH \
     --which_feature $LOST_FEATURES \
diff --git a/scripts/run-single-image.sh b/scripts/run-single-image.sh
index 35888852e55af5d63919b0ed978d1b5088be463d..e9059f2f292ff8376af7eb40a2327d377d452ba3 100755
--- a/scripts/run-single-image.sh
+++ b/scripts/run-single-image.sh
@@ -1,11 +1,8 @@
 declare -a images=(
-    #"000011" "004540" "004541" "007256" "007259" "007267" "007265" "007271"
-    #"000188" "000085" "000038" "000056" "000093" "000263"
-    #"000028" "000048" "000066" "004459" "000435" "003333" "000291"
-    "000003"
+    "COCO_train2014_000000000081" # SOME IMAGE NAME
 )
 
-DATASET_PATH=$KITTI_ROOT/training/image_2/
+DATASET_PATH=$COCO_ROOT/train2014/
 LOST_PATH=~/lost
 DINO_PATH=$LOST_PATH/dino
 OUTPUT_PATH=$LOST_PATH/outputs/samples
@@ -20,12 +17,12 @@ rm -rf $OUTPUT_PATH
 for i in "${images[@]}"
 do
     echo evaluating $i
-    mkdir -p $OUTPUT_PATH/$i/lost
+    mkdir -p $OUTPUT_PATH/$i
     cd $LOST_PATH
     
     python main_lost.py \
-        --image_path $DATASET_PATH/$i.png \
-        --output_dir $OUTPUT_PATH/$i/lost \
+        --image_path $DATASET_PATH/$i.jpg \
+        --output_dir $OUTPUT_PATH/$i \
         --arch $DINO_ARCH \
         --which_feature $LOST_FEATURES \
         --k_patches $K_PATCHES \
@@ -33,8 +30,8 @@ do
         --num_init_seeds 1
 
     python main_lost.py \
-        --image_path $DATASET_PATH/$i.png \
-        --output_dir $OUTPUT_PATH/$i/lost \
+        --image_path $DATASET_PATH/$i.jpg \
+        --output_dir $OUTPUT_PATH/$i \
         --arch $DINO_ARCH \
         --which_feature $LOST_FEATURES \
         --k_patches $K_PATCHES \
@@ -42,8 +39,8 @@ do
         --num_init_seeds 1
 
     python main_lost.py \
-        --image_path $DATASET_PATH/$i.png \
-        --output_dir $OUTPUT_PATH/$i/lost \
+        --image_path $DATASET_PATH/$i.jpg \
+        --output_dir $OUTPUT_PATH/$i \
         --arch $DINO_ARCH \
         --which_feature $LOST_FEATURES \
         --k_patches $K_PATCHES \