diff options
| author | Jordan Gong <jordan.gong@protonmail.com> | 2021-02-15 12:03:07 +0800 |
|---|---|---|
| committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-02-15 12:03:07 +0800 |
| commit | cf59a1f8339e937e31281907a61f484af99ce955 (patch) | |
| tree | aa72b2091d36562c31237c63dd2b0629359d7fb1 /models/rgb_part_net.py | |
| parent | d51312415a32686793d3f0d14eda7fa7cc3990ea (diff) | |
Add DataParallel support on new codebase
Diffstat (limited to 'models/rgb_part_net.py')
| -rw-r--r-- | models/rgb_part_net.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py index 2aa680c..66609fd 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -83,7 +83,7 @@ class RGBPartNet(nn.Module): 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)) + losses = (*losses, hpm_ba_trip, pn_ba_trip) return losses, images else: return x.unsqueeze(1).view(-1) |