diff --git a/model/networks.py b/model/networks.py
index ea6d8ab..a40c527 100644
--- a/model/networks.py
+++ b/model/networks.py
@@ -193,7 +193,7 @@ class DispNetS(TimedModule):
         return torch.nn.Sequential(
             torch.nn.Conv2d(in_planes, out_planes, kernel_size=3, padding=1),
             # TODO try this
-            torch.nn.LayerNorm(out_planes),
+            torch.nn.InstanceNorm2d(out_planes),
             torch.nn.ReLU(inplace=True)
         )