Update `app.py`

app.py

@@ -91,6 +91,7 @@ def load_image_embeddings():
 
 # Use the embeddings in your app
 image_embeddings = load_image_embeddings()
+image_embeddings = image_embeddings.to(DEVICE)
 image_embeddings.shape
 print(image_embeddings.shape)
 
@@ -147,7 +148,7 @@ def get_similar_images(caption, brightness, num_items):
     )
     top_indices = similarities.argsort(descending=True)[:1]
     print(top_indices, similarities[top_indices])
-    return [dataset[i] for i in top_indices.cpu().numpy()]
+    return [dataset[int(i)] for i in top_indices.cpu().numpy()]
 
 
 # Gradio interface