add LayerNorm to conv() for testing

2021-11-29 17:54:45 +01:00 · 2021-11-29 17:54:45 +01:00 · 7060bd5555
commit 7060bd5555
parent 88e4a7e317
1 changed files with 2 additions and 0 deletions
--- a/model/networks.py
+++ b/model/networks.py
@ -192,6 +192,8 @@ class DispNetS(TimedModule):
    def conv(self, in_planes, out_planes):
        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.ReLU(inplace=True)
        )