Fix dropout layer

Author: yuki-inaho

segmentation_models_pytorch/base/heads.py

@@ -5,11 +5,15 @@
 class SegmentationHead(nn.Sequential):
 
     def __init__(self, in_channels, out_channels, kernel_size=3, dropout=None, activation=None, upsampling=1):
-        dropout = nn.Dropout(p=dropout, inplace=False) if dropout else nn.Identity()
+        if dropout:
+            dropout = nn.Dropout(p=dropout, inplace=False) if dropout else nn.Identity()
         conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2)
         upsampling = nn.UpsamplingBilinear2d(scale_factor=upsampling) if upsampling > 1 else nn.Identity()
         activation = Activation(activation)
-        super().__init__(dropout, conv2d, upsampling, activation)
+        if dropout:
+            super().__init__(dropout, conv2d, upsampling, activation)
+        else:
+            super().__init__(conv2d, upsampling, activation)
 
 
 class ClassificationHead(nn.Sequential):