BLIP: enable device map (#34850)

fix device map
huggingface · Dec 11, 2024 · 9094b87 · 9094b87
1 parent 10feacd
commit 9094b87
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/models/blip/modeling_blip.py b/src/transformers/models/blip/modeling_blip.py
@@ -464,6 +464,7 @@ class BlipPreTrainedModel(PreTrainedModel):
     config_class = BlipConfig
     base_model_prefix = "blip"
     supports_gradient_checkpointing = True
+    _no_split_modules = ["BlipEncoderLayer"]
 
     def _init_weights(self, module):
         """Initialize the weights"""