Update README.md

README.md

@@ -16,7 +16,8 @@ examples = ['My text 1', 'This is another text']
 def extract_embeddings(texts):
   encoded_texts = tokenizer(texts)
   with torch.no_grad():
-    style_embeddings = model(encoded_texts.input_ids, attention_mask=encoded_texts.attention_mask)
+    style_embeddings = model(encoded_texts.input_ids,
+                             attention_mask=encoded_texts.attention_mask).pooler_output
   return style_embeddings
 
 print(extract_embeddings(examples))