Merge pull request neuraloperator#183 from btolooshams/main

MLP additional statement to check for scenario where n_layers=1
ziqi-ma · Jul 22, 2023 · 7f9ec18 · 7f9ec18
2 parents e4b9ae6 + b369d7c
commit 7f9ec18
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/neuralop/models/mlp.py b/neuralop/models/mlp.py
@@ -36,7 +36,9 @@ def __init__(self, in_channels, out_channels=None, hidden_channels=None,
         Conv = getattr(nn, f'Conv{n_dim}d')
         self.fcs = nn.ModuleList()
         for i in range(n_layers):
-            if i == 0:
+            if i == 0 and i == (n_layers - 1):
+                self.fcs.append(Conv(self.in_channels, self.out_channels, 1))
+            elif i == 0:
                 self.fcs.append(Conv(self.in_channels, self.hidden_channels, 1))
             elif i == (n_layers - 1):
                 self.fcs.append(Conv(self.hidden_channels, self.out_channels, 1))