Network Training is too slow for Custom Network even though network is not too large(Resnet + LSTM type network)

I have made a custom network for the image data. But the training on this network is very slow though the network is not too huge.

When I am training on resnet150 with lstm, the training is quite fast. Not sure which operation is lagging my training speed?

I have tried reducing Batch_size, seq_dim and network parameters. My guess is some operation is reducing the speed

The dataset I am using are images and they are also very small size(96*96) gray scale images.

***** CODE *****

class Residual(nn.Module):
    def __init__(self, input_channels, num_channels,
                 use_1x1conv=True, strides=1, dilation=2, padding=1, kernel_size=5):
        super(Residual, self).__init__()
        self.conv1 = nn.Conv2d(input_channels, num_channels,
                               kernel_size=kernel_size, padding=padding, stride=strides, dilation=dilation)
        self.conv2 = nn.Conv2d(num_channels, num_channels,
                               kernel_size=kernel_size, padding=2 * padding, dilation=(2 * dilation))
        if use_1x1conv:
            self.conv3 = nn.Conv2d(input_channels, num_channels,
                                   kernel_size=1, stride=strides)
        else:
            self.conv3 = None
        self.bn1 = nn.BatchNorm2d(num_channels)
        self.bn2 = nn.BatchNorm2d(num_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, X):
        Y = F.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3:
            X = self.conv3(X)
        Y += X
        return F.relu(Y)


def weight_init(m):
    '''
    Usage:
        model = Model()
        model.apply(weight_init)
    '''
    if isinstance(m, nn.Conv1d):
        init.normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    elif isinstance(m, nn.Conv2d):
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    elif isinstance(m, nn.BatchNorm1d):
        init.normal_(m.weight.data, mean=1, std=0.02)
        init.constant_(m.bias.data, 0)
    elif isinstance(m, nn.BatchNorm2d):
        init.normal_(m.weight.data, mean=1, std=0.02)
        init.constant_(m.bias.data, 0)
    elif isinstance(m, nn.Linear):
        init.xavier_normal_(m.weight.data)
        init.normal_(m.bias.data)
    elif isinstance(m, nn.LSTM):
        for param in m.parameters():
            if len(param.shape) >= 2:
                init.orthogonal_(param.data)
            else:
                init.normal_(param.data)


class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        num_channels = 64
        strides = 1
        dilation = 2
        padding = 4
        kernel_size = 5
        input_channel = 3
        self.fc = nn.Linear(config['hidden_dim'], config['output_dim'])
        self.lstm = None

        b1 = Residual(input_channels=input_channel, num_channels=num_channels,
                      use_1x1conv=True, strides=strides, dilation=dilation, padding=padding, kernel_size=kernel_size)

        b2 = Residual(input_channels=num_channels, num_channels=2 * num_channels,
                      use_1x1conv=True, strides=strides, dilation=2 * dilation, padding=2 * padding,
                      kernel_size=kernel_size)

        b3 = Residual(input_channels=2 * num_channels, num_channels=4 * num_channels,
                      use_1x1conv=True, strides=strides, dilation=4 * dilation, padding=4 * padding,
                      kernel_size=kernel_size)

        self.net = nn.Sequential(b1, b2, b3, nn.AdaptiveMaxPool2d((2, 2)))
        self.apply(weight_init)

    def forward(self, x):
        x = self.net(x)
        x = x.view(config['batch_size'], config['seq_dim'], -1)
        if self.lstm is None:
            self.lstm = nn.LSTM(x.size(2), config['hidden_dim'], 1, batch_first=True).to(self.device)
            for param in self.lstm.parameters():
                if len(param.shape) >= 2:
                    init.orthogonal_(param.data)
                else:
                    init.normal_(param.data)

        h0 = torch.zeros(config['layer_dim'], x.size(0), config['hidden_dim']).to(self.device)
        # Initialize cell state
        c0 = torch.zeros(config['layer_dim'], x.size(0), config['hidden_dim']).to(self.device)

        output, (hn, cn) = self.lstm(x, (h0,c0))
        output = output[:, :, :]
        output = self.fc(output)

        return output, (hn, cn)

 for epoch in range(config['num_epochs']):
            print('epoch', epoch)
            running_loss = 0
            nb_classes = config['output_dim']
            confusion_matrix = torch.zeros(nb_classes, nb_classes)
            for i, (image, label) in enumerate(trainLoader):
                print('batch: ',i)
                image = image.float().to(device)
                label = label.to(device)
                optimizer.zero_grad()

                batch_size, timesteps, H, W, C = image.size()
                # Change Image shape
                image = image.view(batch_size * timesteps, H, W, C)
                image = image.permute(0, 3, 1, 2)  # from NHWC to NCHW

                output, (hn,cn) = model(image)
                label = label.view(-1)
                output = output.view(-1, output.size(2))
                loss = criterion(output, label)
                loss *= config['seq_dim']
                loss.backward()  # Backward pass
                optimizer.step()  # Now we can do an optimizer step
                running_loss += loss.item()