image width to column

app.py

@@ -66,7 +66,7 @@ Based on [paper:](https://openreview.net/forum?id=1Fqg133qRaI) *Towards Faster a
         picks=[False]*batch_size
         for j,im in enumerate(ims):
             i=j%col_num
-            cols[i].image(im)
+            cols[i].image(im, use_column_width=True)
             picks[j]=cols[i].button("Find Nearest",key="pick_"+str(j))
         
         if any(picks):
@@ -75,7 +75,7 @@ Based on [paper:](https://openreview.net/forum?id=1Fqg133qRaI) *Towards Faster a
                 if pick:
                     scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(ims[i]), k=5)
                     for r in retrieved_examples["image"]:
-                        cols[i].image(r)
+                        cols[i].image(r, use_column_width=True)
     st.write("Nearest neighbors found in the training set according to L2 distance on 'microsoft/beit-base-patch16-224' embeddings")
     st.write(f"Latent dimension: {model.latent_dim}, image size:{model.image_size}")