Update modeling_cge2.py

modeling_cge2.py

@@ -499,7 +499,7 @@ class CgeForEmbedding(CGEModel):
         embeddings = self.get_sentence_embedding(self.embedding_method, hidden_states, 'disc', attention_mask)
         if not return_dict:
             return (embeddings,)
-        return {"sentence_embeddings": embeddings}
+        return {"sentence_embedding": embeddings}
 
         
     def encode(self, sentences, batch_size=16, convert_to_numpy=False,