Allow device auto map (#8)

- feat: no splt modules for device auto map (f4624e02311a18676ed10090705c09efc0f698c2)

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"""