Forgot to add the code aswell
def get_grad_kernel(channels):
  Iy = [[1,-1]]        #We want to implement these two filters

  Ix = [[1],[-1]]      #This part of code was also taken from another answer which we werent able to implement

  return torch.Tensor([[Ix]*channels,[Iy]*channels])/4

class Net(nn.Module):

    def __init__(self):

        super().__init__()

        self.conv1 = nn.Conv2d(3, 64,5,1,2)

        self.pool = nn.MaxPool2d(2, 2)

        self.conv2 = nn.Conv2d(64, 128,5,1,2)

        self.conv3 = nn.Conv2d(128,128,5,1,2)

        self.conv4 = nn.Conv2d(128, 128,5,1,2)

        self.conv5 = nn.Conv2d(128, 128,5,1,2)

        self.fc1 = nn.Linear(128, 120)

        self.fc2 = nn.Linear(120, 64)

        self.fc3 = nn.Linear(64,32)

        self.fc4 = nn.Linear(32,16)

        self.fc5 = nn.Linear(16, 100)

        weights1 = torch.tensor([[0,-1],])

    def forward(self, x):

        x = self.pool(F.relu(self.conv1(x)))

        x = self.pool(F.relu(self.conv2(x)))

        x = self.pool(F.relu(self.conv3(x)))

        x = self.pool(F.relu(self.conv4(x)))

        x = self.pool(F.relu(self.conv5(x)))

        x = torch.flatten(x, 1)  # flatten all dimensions except batch

        x = F.relu(self.fc1(x))

        x = F.relu(self.fc2(x))

        x = F.relu(self.fc3(x))

        x = F.relu(self.fc4(x))

        x = self.fc5(x)

        return x

net = Net()