Spaces:
Running
on
Zero
Running
on
Zero
ahmed-masry
commited on
Commit
•
21cc2de
1
Parent(s):
851a649
Update app.py
Browse files
app.py
CHANGED
@@ -36,7 +36,7 @@ def search(query: str, ds, images, k):
|
|
36 |
|
37 |
qs = []
|
38 |
with torch.no_grad():
|
39 |
-
batch_query = process_queries(
|
40 |
batch_query = {k: v.to(device) for k, v in batch_query.items()}
|
41 |
embeddings_query = model(**batch_query)
|
42 |
qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
|
@@ -80,7 +80,7 @@ def index_gpu(images, ds):
|
|
80 |
images,
|
81 |
batch_size=4,
|
82 |
shuffle=False,
|
83 |
-
collate_fn=
|
84 |
)
|
85 |
|
86 |
|
|
|
36 |
|
37 |
qs = []
|
38 |
with torch.no_grad():
|
39 |
+
batch_query = processor.process_queries([query])
|
40 |
batch_query = {k: v.to(device) for k, v in batch_query.items()}
|
41 |
embeddings_query = model(**batch_query)
|
42 |
qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
|
|
|
80 |
images,
|
81 |
batch_size=4,
|
82 |
shuffle=False,
|
83 |
+
collate_fn=processor.process_images,
|
84 |
)
|
85 |
|
86 |
|