init flamingo embeds new weights

mlfoundations · Sep 17, 2023 · 4875822 · 4875822
1 parent 82d1c69
commit 4875822
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/open_flamingo/src/vlm.py b/open_flamingo/src/vlm.py
@@ -64,6 +64,10 @@ def __init__(
             _weight=self.lang_model.get_input_embeddings().weight,
             pad_token_id=self.pad_token_id,
         )
+        input_embeds.additional_fc.weight.data.normal_(
+            mean=0.0, std=self.lang_model.config.initializer_range
+        )
+
         self.lang_model.set_input_embeddings(input_embeds)
 
         out_embeds = DecoupledLinear(
@@ -72,6 +76,9 @@ def __init__(
             _weight=self.lang_model.get_output_embeddings().weight,
             _bias=self.lang_model.get_output_embeddings().bias,
         )
+        out_embeds.additional_fc.weight.data.normal_(
+            mean=0.0, std=self.lang_model.config.initializer_range
+        )
         self.lang_model.set_output_embeddings(out_embeds)
 
         # gradient checkpointing