Update cognitivess_model/modeling_cognitivess.py

cognitivess_model/modeling_cognitivess.py

@@ -74,7 +74,8 @@ class TransformerBlock(nn.Module):
         return hidden_states
 
 class CognitivessForCausalLM(PreTrainedModel):
-    config_class = CognitivessConfig  # Ensure this matches your config class
+    config_class = CognitivessConfig
+    _no_split_modules = []  # This line ensures that device_map='auto' works
 
     def __init__(self, config: CognitivessConfig):
         super().__init__(config)