Update README.md

README.md

@@ -24,11 +24,8 @@ class MambaTransformerSimple(nn.Module):
             activation="relu",
             batch_first=False,
         )
-        # self.transformer_encoder = nn.TransformerEncoder(
-        #     self.transformer_encoder_layer, num_layers=num_layers
-        # )
         self.transformer_encoder = nn.TransformerEncoder(
-            self.transformer_encoder_layer, num_layers=2
+            self.transformer_encoder_layer, num_layers=num_layers
         )
         self.input_proj = nn.Linear(d_feat, hidden_size)
         self.mamba = Mamba(