Prepare for DataParallel

2 files changed, 11 insertions, 3 deletions

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()
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index 8ebcfd3..2aa680c 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -44,7 +44,8 @@ class RGBPartNet(nn.Module):
             ae_feature_channels * 2, out_channels, hpm_use_1x1conv,
             hpm_scales, hpm_use_avg_pool, hpm_use_max_pool
         )
-        empty_fc = torch.empty(self.hpm_num_parts + tfa_num_parts,
+        self.num_total_parts = self.hpm_num_parts + tfa_num_parts
+        empty_fc = torch.empty(self.num_total_parts,
                                out_channels, embedding_dims)
         self.fc_mat = nn.Parameter(empty_fc)
 
@@ -75,8 +76,13 @@ class RGBPartNet(nn.Module):
         x = self.fc(x)
 
         if self.training:
-            hpm_ba_trip = self.hpm_ba_trip(x[:self.hpm_num_parts], y)
-            pn_ba_trip = self.pn_ba_trip(x[self.hpm_num_parts:], y)
+            y = y.T
+            hpm_ba_trip = self.hpm_ba_trip(
+                x[:self.hpm_num_parts], y[:self.hpm_num_parts]
+            )
+            pn_ba_trip = self.pn_ba_trip(
+                x[self.hpm_num_parts:], y[self.hpm_num_parts:]
+            )
             losses = torch.stack((*losses, hpm_ba_trip, pn_ba_trip))
             return losses, images
         else: