diff options
| author | Jordan Gong <jordan.gong@protonmail.com> | 2021-02-14 16:46:56 +0800 |
|---|---|---|
| committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-02-14 16:46:56 +0800 |
| commit | 929c48093c9f49a515420eb28d2678e48756b300 (patch) | |
| tree | 9ec187d721c97a588f0207efe1311ceeee827d96 /models/model.py | |
| parent | ec7d85a066592bfbed2f6c4c56f9f0b6269897e3 (diff) | |
Prepare for DataParallel
Diffstat (limited to 'models/model.py')
| -rw-r--r-- | models/model.py | 2 |
1 files changed, 2 insertions, 0 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() |