From a31eb135c8cc3a8737fabca54fe3d5791f293753 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Fri, 19 Feb 2021 21:03:17 +0800
Subject: Correct cross reconstruction loss calculated in DataParallel

---
 models/model.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

(limited to 'models/model.py')

diff --git a/models/model.py b/models/model.py
index 4f24ac5..cac165e 100644
--- a/models/model.py
+++ b/models/model.py
@@ -198,12 +198,8 @@ class Model:
             # Duplicate labels for each part
             y = y.unsqueeze(1).repeat(1, self.rgb_pn.module.num_total_parts)
             losses, images = self.rgb_pn(x_c1, x_c2, y)
-            losses = torch.stack((
-                # xrecon           cano_cons         pose_sim
-                losses[0].sum(), losses[1].mean(), losses[2].mean(),
-                # hpm_ba_trip       pn_ba_trip
-                losses[3].mean(), losses[4].mean()
-            ))
+            # Combine losses from different data splits
+            losses = losses.mean()
             loss = losses.sum()
             loss.backward()
             self.optimizer.step()
-- 
cgit v1.2.3