Spaces:
Sleeping
Sleeping
Update demo_support/retrieval.py
Browse files
demo_support/retrieval.py
CHANGED
@@ -38,7 +38,8 @@ feats = deser['feats']
|
|
38 |
def retrieve(embedding, top, sim_th=0.0, filter_fn=None):
|
39 |
sims = []
|
40 |
embedding = F.normalize(embedding.detach().cpu(), dim=-1).squeeze()
|
41 |
-
for chunk in torch.split(feats, 10240):
|
|
|
42 |
sims.append(embedding @ F.normalize(chunk.float(), dim=-1).T)
|
43 |
sims = torch.cat(sims)
|
44 |
sims, idx = torch.sort(sims, descending=True)
|
|
|
38 |
def retrieve(embedding, top, sim_th=0.0, filter_fn=None):
|
39 |
sims = []
|
40 |
embedding = F.normalize(embedding.detach().cpu(), dim=-1).squeeze()
|
41 |
+
#for chunk in torch.split(feats, 10240):
|
42 |
+
for chunk in feats:
|
43 |
sims.append(embedding @ F.normalize(chunk.float(), dim=-1).T)
|
44 |
sims = torch.cat(sims)
|
45 |
sims, idx = torch.sort(sims, descending=True)
|