diff --git a/main_lost.py b/main_lost.py
index 63cb433027714bb423067abc82b1ab21e463aea1..a4eadb5e8a9fc82112cdd65de1550c033819d1c8 100755
--- a/main_lost.py
+++ b/main_lost.py
@@ -298,6 +298,9 @@ if __name__ == "__main__":
                 image = dataset.load_image(im_name)
                 visualize_predictions(image, pred, seed, scales, [w_featmap, h_featmap], vis_folder, im_name)
 
+        # Save the prediction
+        preds_dict[im_name] = pred
+
         # Evaluation
         if args.no_evaluation:
             continue