trying skip connection

models/base.py

@@ -53,13 +53,13 @@ class AEEncoder(nn.Module):
         else:
             x = inputs
 
-        x = self.encoder(x)
+        x, features = self.encoder(x)
         x = self.act(x)
 
         if len(inputs) == 2:
-            return x, inputs[1]
+            return x, features, inputs[1]
         else:
-            return x
+            return x, features
 
 
 class Decoder(nn.Module):
@@ -75,12 +75,12 @@ class Decoder(nn.Module):
             self.softmax = nn.Softmax(dim=1)
 
     def forward(self, inputs):
-        if len(inputs) == 2:
-            x, mask = inputs
+        if len(inputs) == 3:
+            x, features, mask = inputs
         else:
-            x = inputs
+            x, features = inputs
 
-        x = self.decoder(x)
+        x = self.decoder([x, features])
 
         if self.is_label:
             x = x * mask[:,None,:,:]