fix MLPBlock hidden_dim

ChEB-AI · Feb 4, 2025 · ad24fa7 · ad24fa7
1 parent b7ca0e5
commit ad24fa7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chebai/models/ffn.py b/chebai/models/ffn.py
@@ -25,7 +25,7 @@ def __init__(
         current_layer_input_size = input_size
         for hidden_dim in hidden_layers:
             layers.append(MLPBlock(current_layer_input_size, hidden_dim))
-            layers.append(Residual(MLPBlock(current_layer_input_size, hidden_dim)))
+            layers.append(Residual(MLPBlock(hidden_dim, hidden_dim)))
             current_layer_input_size = hidden_dim
 
         layers.append(torch.nn.Linear(current_layer_input_size, self.out_dim))