From de911a563fc503114559d7e0e7f710db090cec0d Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Sat, 9 Jan 2021 21:54:10 +0800
Subject: Add prototype predict function

---
 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 456695d..f39b40b 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -84,7 +84,7 @@ class RGBPartNet(nn.Module):
             loss = torch.sum(torch.stack(losses))
             return loss, [loss.item() for loss in losses]
         else:
-            return x
+            return x.unsqueeze(1).view(-1)
 
     def _disentangle(self, x_c1, x_c2=None, y=None):
         num_frames = len(x_c1)
-- 
cgit v1.2.3