Simple custom activation function causes CUDA out of memory

Soka_David: <h1>Simple custom activation function causes CUDA out of memory.</h1>
<h2>Tags: Custom Activation Function, Memory Usage, CUDA out of memory</h2>
<h2>Problem</h2>
<p>I am trying to implement a very simple activation function that turns every value that is higher than 1 to 1. I did the following:</p>
<pre><code class="lang-python">def zolu(input):
    input[input&gt;1] = 1
    return input

class ZOLU(nn.Module):

    def __init__(self):
        super().__init__() # init the base class

    def forward(self, input):
        return zolu(input)
</code></pre>
<p>And used the function as an activation in my also rather simple model:</p>
<pre><code class="lang-python">class SuperResNet(nn.Module):

    def __init__(self):
        self.size = MODEL_INPUT_SIZE
        super(SuperResNet, self).__init__()
        self.up = nn.Upsample(size=MODEL_OUTPUT_SIZE[1:], mode="bicubic", align_corners=False)
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=5, padding=2)
        self.conv2 = nn.Conv2d(in_channels=10, out_channels=50, kernel_size=5, padding=2)
        self.conv3 = nn.Conv2d(in_channels=50, out_channels=10, kernel_size=5, padding=2)
        self.conv4 = nn.Conv2d(in_channels=10, out_channels=3, kernel_size=5, padding=2)

    def forward(self, x):
        x = zolu(self.up(x))
        x = zolu(self.conv1(x))
        # x = TF.relu(x)
        x = zolu(self.conv2(x))
        # x = TF.relu(x)
        x = zolu(self.conv3(x))
        # x = TF.relu(x)
        x = zolu(self.conv4(x))
        # x = TF.relu(x)
        return x
</code></pre>
<p>When i try to run training i get this output:</p>
<pre><code class="lang-python">Traceback (most recent call last):
  File "C:/Users/sokad/PycharmProjects/SuperResoStrekal/scripts/train.py", line 59, in &lt;module&gt;
    outputs = super_res(crappy_train)
  File "C:\Python37\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Users\sokad\PycharmProjects\SuperResoStrekal\scripts\model.py", line 47, in forward
    x = zolu(self.conv2(x))
  File "C:\Users\sokad\PycharmProjects\SuperResoStrekal\scripts\model.py", line 19, in zolu
    input[input&gt;1] = 1
RuntimeError: CUDA out of memory. Tried to allocate 5.62 GiB (GPU 0; 8.00 GiB total capacity; 1.16 GiB already allocated; 4.48 GiB free; 1.57 GiB reserved in total by PyTorch)
</code></pre>
<h1>Further Infos</h1>
<ul>
<li>instead of input[input &gt; 1] i also tried input[input.gt(1)]</li>
<li>when i switch the activations to RELU, everything works fine</li>
<li>Total params: 26,573</li>
<li>Trainable params: 26,573</li>
<li>Input size (MB): 4.79</li>
<li>Forward/backward pass size (MB): 546.48</li>
<li>Input shape: torch.Size([4, 3, 560, 748])</li>
<li>Batch-Size = 4</li>
<li>Using Sampler-Class (CUDA was OOM without a sampler)</li>
</ul>

ptrblck: <p>You could try to replace the indexing of <code>input</code> with <code>input = torch.clamp(input, max=1)</code>, which might reduce the memory footprint.<br>
Let me know, if that helps.</p>

Soka_David: <p>Thank you very much! It works totally fine! I just didn’t know the clamp()-function yet.</p>