fix SCEModule.sSE (#113)

qubvel-org · Dec 3, 2019 · 9ebf489 · 9ebf489
1 parent a242d53
commit 9ebf489
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/segmentation_models_pytorch/base/modules.py b/segmentation_models_pytorch/base/modules.py
@@ -58,7 +58,7 @@ def __init__(self, in_channels, reduction=16):
             nn.Conv2d(in_channels // reduction, in_channels, 1),
             nn.Sigmoid(),
         )
-        self.sSE = nn.Sequential(nn.Conv2d(in_channels, in_channels, 1), nn.Sigmoid())
+        self.sSE = nn.Sequential(nn.Conv2d(in_channels, 1, 1), nn.Sigmoid())
 
     def forward(self, x):
         return x * self.cSE(x) + x * self.sSE(x)