fixed out of mem: relu needs inplace=True

dicarlolab · Oct 12, 2018 · 37fc7c3 · 37fc7c3
1 parent 1f47e41
commit 37fc7c3
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/cornet_r.py b/cornet_r.py
@@ -27,12 +27,12 @@ def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, w=None):
         self.conv_input = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,
                                     stride=stride, padding=kernel_size // 2)
         self.norm_input = nn.GroupNorm(32, out_channels)
-        self.nonlin_input = nn.ReLU()
+        self.nonlin_input = nn.ReLU(inplace=True)
 
         self.conv1 = nn.Conv2d(out_channels, out_channels,
                                kernel_size=3, padding=1, bias=False)
         self.norm1 = nn.GroupNorm(32, out_channels)
-        self.nonlin1 = nn.ReLU()
+        self.nonlin1 = nn.ReLU(inplace=True)
 
         self.output = Identity()
 

diff --git a/cornet_s.py b/cornet_s.py
@@ -31,15 +31,15 @@ def __init__(self, in_channels, out_channels, times=1):
 
         self.conv1 = nn.Conv2d(out_channels, out_channels * self.scale,
                                kernel_size=1, bias=False)
-        self.nonlin1 = nn.ReLU()
+        self.nonlin1 = nn.ReLU(inplace=True)
 
         self.conv2 = nn.Conv2d(out_channels * self.scale, out_channels * self.scale,
                                kernel_size=3, stride=2, padding=1, bias=False)
-        self.nonlin2 = nn.ReLU()
+        self.nonlin2 = nn.ReLU(inplace=True)
 
         self.conv3 = nn.Conv2d(out_channels * self.scale, out_channels,
                                kernel_size=1, bias=False)
-        self.nonlin3 = nn.ReLU()
+        self.nonlin3 = nn.ReLU(inplace=True)
 
         self.output = Identity()
 
@@ -83,12 +83,12 @@ def CORnet_S():
             ('conv1', nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                             bias=False)),
             ('norm1', nn.BatchNorm2d(64)),
-            ('nonlin1', nn.ReLU()),
+            ('nonlin1', nn.ReLU(inplace=True)),
             ('pool', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),
             ('conv2', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1,
                             bias=False)),
             ('norm2', nn.BatchNorm2d(64)),
-            ('nonlin2', nn.ReLU()),
+            ('nonlin2', nn.ReLU(inplace=True)),
             ('output', Identity())
         ]))),
         ('V2', CORblock_S(64, 128, times=2)),

diff --git a/cornet_z.py b/cornet_z.py
@@ -20,7 +20,7 @@ def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
         super().__init__()
         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,
                               stride=stride, padding=kernel_size // 2)
-        self.nonlin = nn.ReLU()
+        self.nonlin = nn.ReLU(inplace=True)
         self.pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
         self.output = Identity()