From da922be042d96338a3f207386e410b6746d046f5 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Sun, 14 Mar 2021 21:07:28 +0800
Subject: Bug fix when transforming and new config

---
 config.py              | 6 +++---
 models/rgb_part_net.py | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/config.py b/config.py
index 8abeba3..c928067 100644
--- a/config.py
+++ b/config.py
@@ -94,9 +94,9 @@ config: Configuration = {
             'final_gamma': 0.01,
 
             # Local parameters (override global ones)
-            'hpm': {
-                'final_gamma': 0.001
-            }
+            # 'hpm': {
+            #     'final_gamma': 0.001
+            # }
         }
     },
     # Model metadata
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index c38a567..c136040 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -57,7 +57,7 @@ class RGBPartNet(nn.Module):
         if self.training:
             return x_c, x_p, ae_losses, images
         else:
-            return torch.cat((x_c, x_p)).unsqueeze(1).view(-1)
+            return torch.cat((x_c.view(-1), x_p.view(-1)))
 
     def _disentangle(self, x_c1_t2, x_c2_t2=None):
         n, t, c, h, w = x_c1_t2.size()
-- 
cgit v1.2.3