Update app.py

app.py

@@ -21,7 +21,7 @@ def search(query: str, ds, images, k):
     qs = []
     with torch.no_grad():
         batch_query = process_queries(processor, [query], mock_image)
-        batch_query = {k: v.to(device) for k, v in batch_query.items()}
+        batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
         embeddings_query = model(**batch_query)
         qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
 
@@ -55,11 +55,9 @@ def index(files, ds):
         collate_fn=lambda x: process_images(processor, x),
     )
 
-    print(f"model device: {model.device}")
-    print(f"device: {device}")
     for batch_doc in tqdm(dataloader):
         with torch.no_grad():
-            batch_doc = {k: v.to(device) for k, v in batch_doc.items()}
+            batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
             embeddings_doc = model(**batch_doc)
         ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
     return f"Uploaded and converted {len(images)} pages", ds, images