Fix #9 app.py

app.py

@@ -107,7 +107,7 @@ class MappingNetwork(nn.Module):
         return s
 
 class StyleEncoder(nn.Module):
-    def __init__(self, img_size=256, style_dim=64, num_domains=3, max_conv_dim=512):
+    def __init__(self, img_size=256, style_dim=64, num_domains=2, max_conv_dim=512):
         super().__init__()
         dim_in = 64
         blocks = []
@@ -117,20 +117,23 @@ class StyleEncoder(nn.Module):
             dim_out = min(dim_in*2, max_conv_dim)
             blocks += [ResBlk(dim_in, dim_out, downsample=True)]
             dim_in = dim_out
+        blocks += [nn.LeakyReLU(0.2)]
         self.shared = nn.Sequential(*blocks)
+
         self.unshared = nn.ModuleList()
         for _ in range(num_domains):
-            self.unshared += [nn.Linear(dim_in * (img_size // (2**repeat_num))**2, style_dim)]
+            self.unshared += [nn.Linear(dim_in, style_dim)]
 
     def forward(self, x, y):
         h = self.shared(x)
+        h = F.adaptive_avg_pool2d(h, (1,1))
         h = h.view(h.size(0), -1)
         out = []
         for layer in self.unshared:
             out += [layer(h)]
-        out = torch.stack(out, dim=1) # (batch, num_domains, style_dim)
-        idx = torch.LongTensor(range(y.size(0))).unsqueeze(1).to(y.device)
-        s = torch.gather(out, 1, idx.unsqueeze(2).expand(-1, -1, out.size(2))).squeeze(1)
+        out = torch.stack(out, dim=1)
+        idx = torch.arange(y.size(0)).to(y.device)
+        s = out[idx, y]
         return s
 
 # DEFINICIÓN DEL GENERADOR