winfred2027 commited on
Commit
77e656b
·
verified ·
1 Parent(s): 483cfa2

Update demo_support/retrieval.py

Browse files
Files changed (1) hide show
  1. demo_support/retrieval.py +2 -1
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)