Add None activation

Author: qubvel

segmentation_models_pytorch/base/encoder_decoder.py

@@ -10,7 +10,7 @@ def __init__(self, encoder, decoder, activation):
         self.encoder = encoder
         self.decoder = decoder
 
-        if callable(activation):
+        if callable(activation) or activation is None:
             self.activation = activation
         elif activation == 'softmax':
             self.activation = nn.Softmax(dim=1)
@@ -30,6 +30,7 @@ def predict(self, x):
 
         with torch.no_grad():
             x = self.forward(x)
-            x = self.activation(x)
+            if self.activation:
+                x = self.activation(x)
 
         return x