Just add initialization to mlp

config/invariant_mnist.yaml

@@ -20,7 +20,7 @@ mlp :
 
   # Set to true to do polynomial refinement
   p_refine: False
-  normalize: l2 #max_abs # l2 or False
+  normalize: max_abs # l2 or False
 
   input :
     width: 784

examples/invariant_mnist.py

@@ -57,7 +57,7 @@ def __init__(self, cfg):
             normalization=normalization,
         )
 
-        initialize_network_polynomial_layers(self.layer, max_slope=1.0, max_offset=0.0)
+        #initialize_network_polynomial_layers(self.layer, max_slope=1.0, max_offset=0.0)
 
     def setup(self, stage):
         num_train = int(self._train_fraction * 50000)

high_order_layers_torch/networks.py

@@ -180,6 +180,7 @@ def __init__(
         resnet: bool = False,
         device: str = "cpu",
         layer_type_in: str = None,
+        initialization: str="constant_random",
     ) -> None:
         """
         Args :
@@ -207,6 +208,7 @@ def __init__(
             normalization: Normalization to apply after each layer (before any additional nonlinearity).
             resnet: True if layer output should be added to the previous.
             layer_type_in: Layer type for the input layer. If not defined, defaults to layer_type
+            initializtion: layer initialization, "constant_random" or "uniform"
         """
         super().__init__()
         layer_list = []
@@ -224,6 +226,7 @@ def __init__(
             scale=scale,
             periodicity=periodicity,
             device=device,
+            intialization=initialization
         )
         layer_list.append(input_layer)
         for i in range(hidden_layers):
@@ -242,6 +245,7 @@ def __init__(
                 scale=scale,
                 periodicity=periodicity,
                 device=device,
+                initialization=initialization,
             )
 
             # This will add the result of the previous layer after normalization
@@ -264,6 +268,7 @@ def __init__(
             scale=scale,
             periodicity=periodicity,
             device=device,
+            initialization=initialization
         )
         layer_list.append(output_layer)
         self.model = nn.Sequential(*layer_list)