Upload model

config.json

@@ -1,11 +1,10 @@
 {
   "architectures": [
-    "MambaVisionModelForImageClassification"
+    "MambaVisionModel"
   ],
   "auto_map": {
     "AutoConfig": "configuration_mambavision.MambaVisionConfig",
-    "AutoModel": "modeling_mambavision.MambaVisionModel",
-    "AutoModelForImageClassification": "modeling_mambavision.MambaVisionModelForImageClassification"
+    "AutoModel": "modeling_mambavision.MambaVisionModel"
   },
   "depths": [
     3,

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccb5ad877adfa27edfd3bc4d5d8f0b3fd9fdd9be1ec88fa25c1577471ce8201f
+size 390807536

modeling_mambavision.py

@@ -28,7 +28,7 @@ from einops import rearrange, repeat
 
 from transformers import PreTrainedModel
 
-from .configuration_mambavision import MambaVisionConfig
+from configuration_mambavision import MambaVisionConfig
 
 
 def _cfg(url='', **kwargs):
@@ -280,7 +280,7 @@ class ConvBlock(nn.Module):
         self.norm2 = nn.BatchNorm2d(dim, eps=1e-5)
         self.layer_scale = layer_scale
         if layer_scale is not None and type(layer_scale) in [int, float]:
-            self.gamma = nn.Parameter(layer_scale * torch.ones(dim))
+            self.g = nn.Parameter(layer_scale * torch.ones(dim))
             self.layer_scale = True
         else:
             self.layer_scale = False
@@ -294,7 +294,7 @@ class ConvBlock(nn.Module):
         x = self.conv2(x)
         x = self.norm2(x)
         if self.layer_scale:
-            x = x * self.gamma.view(1, -1, 1, 1)
+            x = x * self.g.view(1, -1, 1, 1)
         x = input + self.drop_path(x)
         return x
 
@@ -502,12 +502,12 @@ class Block(nn.Module):
         mlp_hidden_dim = int(dim * mlp_ratio)
         self.mlp = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
         use_layer_scale = layer_scale is not None and type(layer_scale) in [int, float]
-        self.gamma_1 = nn.Parameter(layer_scale * torch.ones(dim))  if use_layer_scale else 1
-        self.gamma_2 = nn.Parameter(layer_scale * torch.ones(dim))  if use_layer_scale else 1
+        self.g_1 = nn.Parameter(layer_scale * torch.ones(dim))  if use_layer_scale else 1
+        self.g_2 = nn.Parameter(layer_scale * torch.ones(dim))  if use_layer_scale else 1
 
     def forward(self, x):
-        x = x + self.drop_path(self.gamma_1 * self.mixer(self.norm1(x)))
-        x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        x = x + self.drop_path(self.g_1 * self.mixer(self.norm1(x)))
+        x = x + self.drop_path(self.g_2 * self.mlp(self.norm2(x)))
         return x