From 901fc638fc83ab7ffd692660dd3a8043b793f887 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Tue, 8 Aug 2023 21:36:04 +0800 Subject: [PATCH] rename variable --- vision_toolbox/backbones/mlp_mixer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vision_toolbox/backbones/mlp_mixer.py b/vision_toolbox/backbones/mlp_mixer.py index b16800d..600c888 100644 --- 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):