From b6e5972b64cc61fc967cf3d098fc629d781adce4 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Mon, 22 Mar 2021 19:32:16 +0800
Subject: Add embedding visualization and validate on testing set

---
 models/rgb_part_net.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'models/rgb_part_net.py')

diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index c136040..4a82da3 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.view(-1), x_p.view(-1)))
+            return x_c, x_p
 
     def _disentangle(self, x_c1_t2, x_c2_t2=None):
         n, t, c, h, w = x_c1_t2.size()
-- 
cgit v1.2.3