Cuda out of memory during evaluation (tried everything)

seanco92: <p>Hi,<br>
I’m facing a cuda OOM error during evaluation only in the prediction phase (model(inputs))<br>
Training loop works fine.<br>
few points:</p>
<ol>
<li>It happens with and without training loop</li>
<li>The train and evaluation loop are in separate functions (each epoch each function called - does not looks like scope issue)</li>
<li>I am using with torch.no_grad() and model.eval()</li>
<li>train and eval loaders has same data types, same shapes and same batch sizes.</li>
<li>This framework works with lot’s of models. The problem occur in new model I implemented - so I suspect that the problem is in the model and not in the code around it</li>
<li>The problem occur with and without mixed precision, but if I use torch.cuda.amp.autocast() before evaluation - the problem is solved. (Still I want to know why)</li>
<li>If I use torch.cuda.clear_cash() inside the forward() method - problem is allegedly solved.</li>
</ol>
<p>My model has 5 layers from the same kind and the input grows becaause of number of filters of convolutional layers.<br>
initial input size is (2048, 10, 256) == (B, C, N) and<br>
The forward function of my model:</p>
<pre><code>def forward(self, x, info_dict=None):
    self.process_info_dict(info_dict)
    tmp_res = []
    for d, m in self.slice_layers.items():
        if d == "remainder":
            tmp_res.append(m(x[:, self.slice_dict[d]]))
        else:
            b, c, n = (x.shape[0], len(self.slice_dict[d]), x.shape[2])
            tmp_x = x[:, self.slice_dict[d]].view(-1, 1, n)
            res = m(tmp_x)
            res = res.view(b, c * self.n_filters, self.output_dim[-1])
            tmp_res.append(res)
    x = torch.cat(tmp_res, 1)
    self.update_info_dict(info_dict)
    return x
</code></pre>
<p>When info_dcit is a dictionary of few hundreds of integers<br>
The error occur in the line<br>
x = torch.cat(tmp_res, 1)<br>
The for loop is executed only 3 times.<br>
The model layers are defined here:</p>
<pre><code class="lang-auto">def create_layers(self, out_h, dtype_slices, kernel_size, dilation, padding, stride, groups):
        for i, (d, layer_out_h) in enumerate(out_h.items()):
            block_layers = []
            in_channels = dtype_slices[d] if d == "remainder" else 1
            out_channels = layer_out_h if d == "remainder" else int(layer_out_h / dtype_slices[d])
            layer = torch.nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, 
                                                   dilation=dilation,
                                                  padding=padding, stride=stride, groups=groups[d])
            block_layers.append(layer)
            if self.activation:
                block_layers.append(self.activation)
            self.output_dim = conv_output_shape(self.output_dim, kernel_size, stride, padding)[0]
            self.slice_layers[d] = (torch.nn.Sequential(*block_layers))
</code></pre>
<p>and all the layers save in:<br>
self.slice_layers = torch.nn.ModuleDict()</p>
<p>error message:<br>
CUDA out of memory. Tried to allocate 5.00 GiB (GPU 0; 21.99 GiB total capacity; 10.30 GiB already allocated; 2.29 GiB free; 19.40 GiB reserved in total by PyTorch) If reserved memory is &gt;&gt; allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF</p>
<p>Any idea what to check and how?<br>
Thanks!</p>