Update init

Author: qubvel

segmentation_models_pytorch/base/initialization.py

@@ -8,7 +8,9 @@ def initialize_decoder(module):
             if m.bias is not None:
                 nn.init.constant_(m.bias, 0)
 
-        elif isinstance(m, nn.BatchNorm2d):
+        elif isinstance(
+            m, (nn.BatchNorm2d, nn.LayerNorm, nn.GroupNorm, nn.InstanceNorm2d)
+        ):
             nn.init.constant_(m.weight, 1)
             nn.init.constant_(m.bias, 0)