Change topk output for I2I mode

app.py

@@ -52,9 +52,9 @@ def image_2_image(model, img_emb, img_names, img_urls,n_top_k_images):
             st.write("The image with the most similar embedding is:")
             cosine_sim = get_match(model, image.convert("RGB"), img_emb)
             logger.info(cosine_sim.shape)
-            top_k_images_indices = torch.topk(cosine_sim, n_top_k_images, 1).indices
+            top_k_images_indices = torch.topk(cosine_sim, n_top_k_images, 0).indices
             logger.info(top_k_images_indices.squeeze().tolist())
-            images_found = [img_names[top_k_best_image] for top_k_best_image in top_k_images_indices.squeeze().tolist()]
+            images_found = [img_names[top_k_best_image] for top_k_best_image in top_k_images_indices]
             cols = st.columns(n_top_k_images)
             for i, image_found in enumerate(images_found):
                 logger.success(f"Image match found: {image_found}")