Upload model

config.json

@@ -1,6 +1,10 @@
 {
+  "architectures": [
+    "MambaModelForCausalLM"
+  ],
   "auto_map": {
-    "AutoConfig": "configuration_mamba.MambaConfig"
+    "AutoConfig": "configuration_mamba.MambaConfig",
+    "AutoModelForCausalLM": "modeling_mamba.MambaModelForCausalLM"
   },
   "bias": false,
   "conv_bias": true,
@@ -14,6 +18,7 @@
   "model_type": "mamba",
   "n_layer": 24,
   "pad_vocab_size_multiple": 8,
+  "torch_dtype": "float32",
   "transformers_version": "4.37.2",
   "vocab_size": 50280
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:287cad4048030ae246aeda26c0e703b838c50422fe89f19099298c034b25e7b5
-size 516565384
+oid sha256:699ed6f59fb948186f449c5031e0dc659d504c90d7e018302aa1e190cdb40220
+size 516567560

modeling_mamba.py

@@ -227,13 +227,13 @@ class MambaModel(MambaPreTrainedModel):
             config: MambaConfig
         """
         super().__init__(config)
-        self.config = config
+        # self.config = config
 
-        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.embedding = nn.Embedding(self.config.vocab_size, self.config.d_model)
         self.layers = nn.ModuleList(
-            [MambaBlock(config, layer_idx) for layer_idx in range(config.n_layer)]
+            [MambaBlock(self.config, layer_idx) for layer_idx in range(self.config.n_layer)]
         )
-        self.norm_f = MambaRMSNorm(config.d_model)
+        self.norm_f = MambaRMSNorm(self.config.d_model)
 
         self.gradient_checkpointing = False
         self.post_init()