From 09c2af6f4881bb4a2ed5e3685b8ce65ad41695cb Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Tue, 2 Mar 2021 20:10:14 +0800
Subject: Fix DataParallel specific bugs

---
 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 8a0f3a7..cdf579b 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -67,7 +67,7 @@ class RGBPartNet(nn.Module):
         x = self.fc(x)
 
         if self.training:
-            return x, ae_losses, images
+            return x.transpose(0, 1), ae_losses, images
         else:
             return x.unsqueeze(1).view(-1)
 
-- 
cgit v1.2.3