Hello,
I have added a few layers to the resnet18 model but when I try to find the summary of the model using a dummy input, I am getting a RuntimeError.
Code:

import torch.nn as nn
import torch
from torchsummary import summary
import torchvision


class BasicBlock2(nn.Module):

    def __init__(self, in_channels, out_channels, stride=(1, 1), dilation=(1, 1), downsample=None):
        super(BasicBlock2, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(5, 5),
                               padding=(0, 0),
                               dilation=dilation,
                               stride=stride)
        self.batch_norm1 = nn.BatchNorm2d(out_channels)
        self.relu1 = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(3, 3),
                               padding=(1, 1),
                               dilation=dilation,
                               stride=stride)
        self.batch_norm2 = nn.BatchNorm2d(out_channels)
        self.relu2 = nn.ReLU(inplace=True)

    def forward(self, x):
        # The input shape of the first layer should be: [-1, 128, 64, 192]
        out = self.conv1(x)
        out = self.batch_norm1(out)
        out = self.relu1(out)

        # out = self.conv2(out)
        # out = self.batch_norm2(out)
        # out = self.relu2(out)
        # The output shape of the first layer should be: [-1, 128, 44, 172]
        return out


class BasicBlock3(nn.Module):

    def __init__(self, in_channels, out_channels, stride=(1, 1), dilation=(1, 1), downsample=None):
        super(BasicBlock3, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(7, 7),
                               padding=(3, 0),
                               dilation=dilation,
                               stride=stride)
        self.batch_norm1 = nn.BatchNorm2d(out_channels)
        self.relu1 = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(3, 3),
                               padding=(1, 1),
                               dilation=dilation,
                               stride=stride)
        self.batch_norm2 = nn.BatchNorm2d(out_channels)
        self.relu2 = nn.ReLU(inplace=True)

    def forward(self, x):
        # The input shape of the first layer should be: [-1, 128, 44, 172]
        out = self.conv1(x)
        out = self.batch_norm1(out)
        out = self.relu1(out)

        # out = self.conv2(out)
        # out = self.batch_norm2(out)
        # out = self.relu2(out)
        # The output shape of the first layer should be: [-1, 128, 44, 124]
        return out


class NewLayers(nn.Module):

    def __init__(self, layers, last_in_channels, last_out_channels, stride=(1, 1), dilation=(1, 1)):
        super(NewLayers, self).__init__()
        # layers: [5, 8]
        # last_in_channels: 128
        # last_out_channels: 256

        # Need two _make_layers() with blocks 5 and 8 respectively
        self.intermediate1_layers = self._make_layer(BasicBlock2, layers[0], in_channels=128, out_channels=128)
        self.intermediate2_layers = self._make_layer(BasicBlock3, layers[1], in_channels=128, out_channels=128)

        self.lastconv = nn.Conv2d(in_channels=last_in_channels,
                                  out_channels=last_out_channels,
                                  kernel_size=(5, 5),
                                  padding=(2, 0),
                                  dilation=dilation,
                                  stride=stride)
        self.lastbatchnorm = nn.BatchNorm2d(last_out_channels)
        self.lastrelu = nn.ReLU(inplace=True)

    def _make_layer(self, block, blocks, in_channels, out_channels):  # blocks = 5 and blocks = 8
        layers = []
        for _ in range(1, blocks+1):  # because in python last index is not considered, hence + 1
            layers.append(block(in_channels, out_channels))
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.intermediate1_layers(x)
        out = self.intermediate2_layers(out)
        out = self.lastconv(out)
        out = self.lastbatchnorm(out)
        out = self.lastrelu(out)
        return out


net = NewLayers(layers=[5, 8], last_in_channels=128, last_out_channels=128)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
x = torch.randn(1, 128, 64, 192)
output= net(x)

resnet = torchvision.models.resnet18()
model = nn.Sequential(*list(resnet.children())[:-4])
model.add_module('custom_layer', net)
model.add_module('final', nn.Sequential(*(list(resnet.children())[-4:])))
print(model)
summary(model.to(device), (3, 512, 1536))

But the error I am getting is:
RuntimeError: mat1 dim 1 must match mat2 dim 0
I cannot find where the dimensions are wrong. Any help is appreciated.
Thank You.