small fix to make model behave the same as before with old weights

podgorskiy · Dec 7, 2020 · 42bcf2e · 42bcf2e
1 parent d43e061
commit 42bcf2e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/net.py b/net.py
@@ -910,7 +910,8 @@ def forward(self, x):
             x = self.map_blocks[i](x)
         # We select just one output. For compatibility with older models.
         # All other outputs are ignored
-        return x[:, 0, 0]
+        # It is the same as if the last layer had one output.
+        return x[:, 0, x.shape[2] // 2]
 
 
 @MAPPINGS.register("MappingDNoStyle")