rename variable

gau-nernst · Aug 8, 2023 · 901fc63 · 901fc63
1 parent ddd8f3e
commit 901fc63
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/vision_toolbox/backbones/mlp_mixer.py b/vision_toolbox/backbones/mlp_mixer.py
@@ -15,11 +15,11 @@
 
 
 class MLP(nn.Sequential):
-    def __init__(self, d_model: int, mlp_dim: float, act: _act = nn.GELU) -> None:
+    def __init__(self, in_dim: int, hidden_dim: float, act: _act = nn.GELU) -> None:
         super().__init__()
-        self.linear1 = nn.Linear(d_model, mlp_dim)
+        self.linear1 = nn.Linear(in_dim, hidden_dim)
         self.act = act()
-        self.linear2 = nn.Linear(mlp_dim, d_model)
+        self.linear2 = nn.Linear(hidden_dim, in_dim)
 
 
 class MixerBlock(nn.Module):