add set/get_output_embeddings

rhymes-ai · Nov 28, 2024 · d7687b6 · d7687b6
1 parent 2814c7c
commit d7687b6
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/aria/model/modeling_aria.py b/aria/model/modeling_aria.py
@@ -164,6 +164,14 @@ def get_input_embeddings(self) -> nn.Module:
     def set_input_embeddings(self, value):
         """Set the input embeddings for the language model."""
         self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        """Retrieve the output embeddings from the language model."""
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, value):
+        """Set the output embeddings for the language model."""
+        self.language_model.set_output_embeddings(value)
 
     def set_moe_z_loss_coeff(self, value):
         """