irow commited on
Commit
b6c11fd
·
1 Parent(s): 04867f7

Update inference.py, for more general use.

Browse files
Files changed (1) hide show
  1. inference.py +2 -2
inference.py CHANGED
@@ -61,8 +61,8 @@ class DualBertModel: # use this class to encode your document and queries.
61
  with torch.no_grad():
62
  to_return = []
63
  for c in self.chunks(corpus, batch_size):
64
- to_encode = [d['text'] for d in c]
65
- ids, atm = self.process(to_encode, False)
66
  out = self.p_model(ids, atm)
67
  to_return.append(out.cpu())
68
 
 
61
  with torch.no_grad():
62
  to_return = []
63
  for c in self.chunks(corpus, batch_size):
64
+ # to_encode = [d['text'] for d in c]
65
+ ids, atm = self.process(c, False)
66
  out = self.p_model(ids, atm)
67
  to_return.append(out.cpu())
68