From 929c48093c9f49a515420eb28d2678e48756b300 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Sun, 14 Feb 2021 16:46:56 +0800
Subject: Prepare for DataParallel

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

(limited to 'models/model.py')

diff --git a/models/model.py b/models/model.py
index 9cb6a8e..f79b832 100644
--- a/models/model.py
+++ b/models/model.py
@@ -193,6 +193,8 @@ class Model:
             x_c1 = batch_c1['clip'].to(self.device)
             x_c2 = batch_c2['clip'].to(self.device)
             y = batch_c1['label'].to(self.device)
+            # Duplicate labels for each part
+            y = y.unsqueeze(1).repeat(1, self.rgb_pn.num_total_parts)
             losses, images = self.rgb_pn(x_c1, x_c2, y)
             loss = losses.sum()
             loss.backward()
-- 
cgit v1.2.3