Update README.md

README.md

@@ -25,12 +25,13 @@ def load_finetuned_model():
         return tokenizer, query_encoder, sentence_encoder
 
 
-def encode_batch(model, tokenizer, sentences, device):
+def encode_batch_fn(model, tokenizer, sentences, device)
     input_ids = tokenizer(sentences, padding=True, max_length=128, truncation=True, return_tensors="pt",
                           add_special_tokens=True).to(device)
     features = model(**input_ids)[0]
-    features =  torch.sum(features[:,1:,:] * input_ids["attention_mask"][:,1:].unsqueeze(-1), dim=1) / torch.clamp(torch.sum(input_ids["attention_mask"][:,1:], dim=1, keepdims=True), min=1e-9)
-        
+
+    features =  torch.sum(features[:,:,:] * input_ids["attention_mask"][:,:].unsqueeze(-1), dim=1) / torch.clamp(torch.sum(input_ids["attention_mask"][:,:], dim=1, keepdims=True), min=1e-9)
+
     return features
 
 ```