And my CNN
class CNN(nn.Module): 
    def __init__(self, in_channels = 3, num_classes = 1):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3, stride = (1,1), padding = (1,1))
        self.pool = nn.MaxPool2d(kernel_size=2, stride = (2,2))
        self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride = (1,1), padding = (1,1))
        self.fc1 = nn.Linear(16*75*75, num_classes)
        self.sigmoid = nn.Sigmoid()


    def forward(self, x):
        x= F.relu(self.conv1(x))
        x = self.pool(x)
        x= F.relu(self.conv2(x))
        x = self.pool(x)
        x = x.reshape(x.shape[0], -1)
        x = self.fc1(x)
        x = self.sigmoid(x)
        return x