Update model/models.py

model/models.py

@@ -71,8 +71,8 @@ class Selayer(nn.Module):
     def __init__(self, inplanes):
         super(Selayer, self).__init__()
         self.global_avgpool = nn.AdaptiveAvgPool2d(1)
-        self.conv1 = nn.Conv2d(inplanes, inplanes // 16, kernel_size=1, stride=1)
-        self.conv2 = nn.Conv2d(inplanes // 16, inplanes, kernel_size=1, stride=1)
+        self.conv1 = nn.Conv2d(2, inplanes // 16, kernel_size=1, stride=1)
+        self.conv2 = nn.Conv2d(inplanes // 16, 2, kernel_size=1, stride=1)
         self.relu = nn.ReLU(inplace=True)
         self.sigmoid = nn.Sigmoid()
 
@@ -89,8 +89,8 @@ class SelayerSpectr(nn.Module):
     def __init__(self, inplanes):
         super(SelayerSpectr, self).__init__()
         self.global_avgpool = nn.AdaptiveAvgPool2d(1)
-        self.conv1 = SpectralNorm(nn.Conv2d(inplanes, inplanes // 16, kernel_size=1, stride=1))
-        self.conv2 = SpectralNorm(nn.Conv2d(inplanes // 16, inplanes, kernel_size=1, stride=1))
+        self.conv1 = SpectralNorm(nn.Conv2d(2, inplanes // 16, kernel_size=1, stride=1))
+        self.conv2 = SpectralNorm(nn.Conv2d(inplanes // 16, 2, kernel_size=1, stride=1))
         self.relu = nn.ReLU(inplace=True)
         self.sigmoid = nn.Sigmoid()