init flamingo embeds new weights

mlfoundations · Sep 17, 2023 · 8f2f040 · 8f2f040
1 parent 4875822
commit 8f2f040
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/open_flamingo/src/vlm.py b/open_flamingo/src/vlm.py
@@ -64,9 +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
-        )
+        if hasattr(input_embeds, "additional_embedding"):
+            input_embeds.additional_embedding.weight.data.normal_(
+                mean=0.0, std=self.lang_model.config.initializer_range
+            )
 
         self.lang_model.set_input_embeddings(input_embeds)
 
@@ -76,9 +77,10 @@ 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
-        )
+        if hasattr(out_embeds, "additional_fc"):
+            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