fix flamingo init

mlfoundations · Mar 20, 2024 · 292afa1 · 292afa1
1 parent 52ca075
commit 292afa1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/open_flamingo/src/vlm.py b/open_flamingo/src/vlm.py
@@ -393,6 +393,7 @@ def __init__(
             gradient_checkpointing=gradient_checkpointing,
         )
         self.lang_model.set_decoder_layers_attr_name(decoder_layers_attr_name)
+        self.decoder_layers_attr_name = decoder_layers_attr_name
         self.lang_model.init_cross_attention_layers(
             lang_hidden_size=self.lang_hidden_dim,
             vis_hidden_size=self.vis_embedding_dim,
@@ -491,7 +492,7 @@ def lambda_fn(module: nn.Module):
                 return True
             if isinstance(module, GatedCrossAttentionBlock):
                 return True
-            if isinstance(module, original_decoder_block_class):
+            if isinstance(module, decoder_block_class):
                 return True
 
         return lambda_fn