diff --git a/main.py b/main.py
index 31f024b..625b6e8 100644
--- a/main.py
+++ b/main.py
@@ -641,7 +641,7 @@ if __name__ == "__main__":
         from skorch import NeuralNet
 
         estimator = NeuralNet(
-            module=FouriER,
+            module=FouriER(model.p),
             criterion=torch.nn.BCELoss,
             optimizer=torch.optim.Adam,
             max_epochs=100,