Update README.md

README.md

@@ -19,7 +19,7 @@ def extract_feature(protein):
   input_ids = torch.tensor(ids['input_ids']).to(self.device)
   attention_mask = torch.tensor(ids['attention_mask']).to(self.device)
   with torch.no_grad():
-    embedding_repr = self.model(output_hidden_states=True, input_ids=input_ids,attention_mask=attention_mask).last_hidden_state
+    embedding_repr = self.model(input_ids=input_ids,attention_mask=attention_mask).last_hidden_state
   return torch.mean(embedding_repr)
 ```
 
@@ -32,7 +32,7 @@ def extract_features_batch(proteins):
   input_ids = torch.tensor(ids['input_ids']).to(self.device)
   attention_mask = torch.tensor(ids['attention_mask']).to(self.device)
   with torch.no_grad():
-    embedding_repr = self.model(output_hidden_states=True, input_ids=input_ids,attention_mask=attention_mask).last_hidden_state
+    embedding_repr = self.model(input_ids=input_ids,attention_mask=attention_mask).last_hidden_state
   attention_mask = attention_mask.unsqueeze(-1)
   attention_mask = attention_mask.expand(-1, -1, embedding_repr.size(-1))
   masked_embedding_repr = embedding_repr * attention_mask