From 42bcf2e5f213ff1c919483678344f3da6bc90f8a Mon Sep 17 00:00:00 2001
From: Stanislav Pidhorskyi <stanislav@podgorskiy.com>
Date: Mon, 7 Dec 2020 08:45:55 -0500
Subject: [PATCH] small fix to make model behave the same as before with old
 weights

---
 net.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/net.py b/net.py
index 09ec3b2b..5a3c4073 100644
--- 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")