import torch 
import torch.nn as  nn 
import torch.functional as F
from bn import batch_norm
class residual(nn.Module):
    def __init__(self, inp, out, stride = 1):
        super().__init__()
        self.bn1 = batch_norm(inp)
        self.conv1 = nn.Conv2d(inp, out, kernel_size=3, padding = 1, stride = stride)
        self.bn2 = batch_norm(out)
        self.conv2 = nn.Conv2d(out, out, kernel_size = 3, padding = 1, stride = 1)
        # skip cpnnection
        self.concat = nn.Conv2d(inp, out, kernel_size = 1, padding = 0, stride = stride)
    def forward(self, input):
        x = self.bn1(input)
        x = self.conv1(x)
        x = self.bn2(x)
        x = self.conv2(x)
        skip = self.concat(input)
        skip = x+skip

        return skip