Built in flatten (#486)

Author: qubvel

segmentation_models_pytorch/base/heads.py

@@ -1,5 +1,5 @@
 import torch.nn as nn
-from .modules import Flatten, Activation
+from .modules import Activation
 
 
 class SegmentationHead(nn.Sequential):
@@ -17,7 +17,7 @@ def __init__(self, in_channels, classes, pooling="avg", dropout=0.2, activation=
         if pooling not in ("max", "avg"):
             raise ValueError("Pooling should be one of ('max', 'avg'), got {}.".format(pooling))
         pool = nn.AdaptiveAvgPool2d(1) if pooling == 'avg' else nn.AdaptiveMaxPool2d(1)
-        flatten = Flatten()
+        flatten = nn.Flatten()
         dropout = nn.Dropout(p=dropout, inplace=True) if dropout else nn.Identity()
         linear = nn.Linear(in_channels, classes, bias=True)
         activation = Activation(activation)

segmentation_models_pytorch/base/modules.py

@@ -118,8 +118,3 @@ def __init__(self, name, **params):
 
     def forward(self, x):
         return self.attention(x)
-
-
-class Flatten(nn.Module):
-    def forward(self, x):
-        return x.view(x.shape[0], -1)