My model size is not so big (2M parameters). However, it throws the following error after several epochs (sometimes it could train for >20 epochs and throw this error).
RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 31.75 GiB total capacity; 16.29 GiB already allocated; 1.69 MiB free; 14.46 GiB cached)
Here is my model:
class WaveNetEncoder9(Module):

    def __init__(self,
                 in_channels: int,
                 out_waveblock: List,
                 kernel_size: int,
                 dilation_rates: List,
                 inner_kernel_size: int,
                 inner_padding: int,
                 last_dim: int,
                 mode: str,
                 ) \
            -> None:
        super(WaveNetEncoder9, self).__init__()

        assert len(dilation_rates) == len(out_waveblock), "length of dilation rates must match out channels"
        self.in_channels: int = in_channels
        self.out_waveblock = out_waveblock
        self.kernel_size = kernel_size
        self.dilation_rates = dilation_rates

        self.wave_blocks = ModuleList() 
        self.bn_relu = ModuleList()
        self.wave_blocks.append(WaveBlock(
            in_channels=self.in_channels,
            out_channels=self.out_waveblock[0],
            dilation_rates=self.dilation_rates[0],
            kernel_size=self.kernel_size
        ))
        
        for i in range(1,len(self.dilation_rates)):
            self.wave_blocks.append(
                WaveBlock(
                    in_channels=self.out_waveblock[i-1],
                    out_channels=self.out_waveblock[i],
                    dilation_rates=self.dilation_rates[i],
                    kernel_size=self.kernel_size
                )
            )
        
        for i in range(len(self.dilation_rates)):
            self.bn_relu.append(
                Sequential(
                    BatchNorm1d(self.out_waveblock[i]),
                    ReLU(),
                )
            )

        if mode == "max":
            self.dnn = DepthWiseSeparableDNN(  # this is 3 blocks of conv2d+relu+batchnorm2d and avgpool2d layers
                cnn_channels=out_waveblock[-1],
                cnn_dropout=0.2,
                inner_kernel_size=inner_kernel_size,
                inner_padding=inner_padding
            )
        elif mode == 'avg':
            self.dnn = DepthWiseSeparableDNN2(
                cnn_channels=out_waveblock[-1],
                cnn_dropout=0.2,
                inner_kernel_size=inner_kernel_size,
                inner_padding=inner_padding
            )
        else:
            raise Exception('Incorrect argument')
        
        self.point_wise: Module = Conv2d(
            in_channels=2, out_channels=1,
            kernel_size=inner_kernel_size, stride=1,
            padding=1, dilation=1,
            groups=1, bias=True, padding_mode='zeros')

        self.fc_audioset = Linear(last_dim,last_dim, bias=True)

        
    def forward(self,
                x: Tensor) \
            -> Tensor:
        """Forward pass of the encoder.

        :param x: Input to the encoder.
        :type x: torch.Tensor
        :return: Output of the encoder.
        :rtype: torch.Tensor
        """
        bs, time_step, mel_bins = x.size()
        x1 = x.permute(0,2,1)
        x2 = x

        # WaveNet
        for i in range(len(self.wave_blocks)):
            x1 = self.wave_blocks[i](x1)
            x1 = self.bn_relu[i](x1) 
        
        x1 = x1.unsqueeze(1).transpose(2,3)
        # CNN
        x2 = self.dnn(x2)
        x2 = x2.transpose(1,3) # (bs, 1, time_steps, 128)
        x2 = x2.reshape(x2.shape[0],1,time_step,-1)
        x = torch.cat((x2, x1), dim=1)
        x = self.point_wise(x) 

        x = x.reshape(bs,time_step,-1)
        x = F.relu_(self.fc_audioset(x))
        return x

# EOF

I also empty cached memory at the begining of each iteration by
torch.cuda.empty_cache()
gc.collect()

use torch.no_grad() in evaluation and detach loss loss.item().