m3hrdadfi commited on
Commit
a890234
1 Parent(s): 275d89d

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -45,7 +45,7 @@ def max_pooling(model_output, attention_mask):
45
  token_embeddings = model_output[0] #First element of model_output contains all token embeddings
46
  input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
47
  token_embeddings[input_mask_expanded == 0] = -1e9 # Set padding tokens to large negative value
48
- return torch.max(token_embeddings, 1)[0]
49
 
50
  # Sentences we want sentence embeddings for
51
  sentences = [
 
45
  token_embeddings = model_output[0] #First element of model_output contains all token embeddings
46
  input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
47
  token_embeddings[input_mask_expanded == 0] = -1e9 # Set padding tokens to large negative value
48
+ return torch.mean(token_embeddings, 1)[0]
49
 
50
  # Sentences we want sentence embeddings for
51
  sentences = [