Allow device auto map

modeling_bert.py

@@ -956,6 +956,7 @@ class JinaBertPreTrainedModel(PreTrainedModel):
     load_tf_weights = load_tf_weights_in_bert
     base_model_prefix = "bert"
     supports_gradient_checkpointing = True
+    _no_split_modules = ["JinaBertLayer"]
 
     def _init_weights(self, module):
         """Initialize the weights"""