Update file model.py

model.py

@@ -1,7 +1,7 @@
 #from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
 #from mamba_ssm.models.config_mamba import MambaConfig
 
-from mamba import Mamba, ModelConfig
+from mamba import Mamba, ModelArgs
 
 
 import torch
@@ -13,7 +13,7 @@ class Model:
     def __init__(self, config: Config):
         self.__dict__ = dict(config.__dict__)
 
-        #self.model = MambaLMHeadModel(MambaConfig(**self.params.__dict__)).to(GetDevice())
+        #self.model = MambaLMHeadModel(ModelArgs(**self.params.__dict__)).to(GetDevice())
         self.model = Mamba(ModelConfig(**self.params.__dict__)).to(GetDevice())
         self.log()