Backend pytorch: FNN supports regularization (#1833)

lululxvi · Sep 7, 2024 · f7aa563 · f7aa563
1 parent d5a903d
commit f7aa563
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/deepxde/nn/pytorch/fnn.py b/deepxde/nn/pytorch/fnn.py
@@ -9,7 +9,9 @@
 class FNN(NN):
     """Fully-connected neural network."""
 
-    def __init__(self, layer_sizes, activation, kernel_initializer):
+    def __init__(
+        self, layer_sizes, activation, kernel_initializer, regularization=None
+    ):
         super().__init__()
         if isinstance(activation, list):
             if not (len(layer_sizes) - 1) == len(activation):
@@ -21,6 +23,7 @@ def __init__(self, layer_sizes, activation, kernel_initializer):
             self.activation = activations.get(activation)
         initializer = initializers.get(kernel_initializer)
         initializer_zero = initializers.get("zeros")
+        self.regularizer = regularization
 
         self.linears = torch.nn.ModuleList()
         for i in range(1, len(layer_sizes)):