diff options
| author | Jordan Gong <jordan.gong@protonmail.com> | 2021-03-02 20:10:14 +0800 |
|---|---|---|
| committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-03-02 20:10:14 +0800 |
| commit | 09c2af6f4881bb4a2ed5e3685b8ce65ad41695cb (patch) | |
| tree | 3e7e9fba6d1fc04f3d976a2aec9a9cc71e712845 /models/rgb_part_net.py | |
| parent | ab09d9bd0fe92d97f340feef1e2bbbcd33468953 (diff) | |
Fix DataParallel specific bugs
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 8a0f3a7..cdf579b 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -67,7 +67,7 @@ class RGBPartNet(nn.Module): x = self.fc(x) if self.training: - return x, ae_losses, images + return x.transpose(0, 1), ae_losses, images else: return x.unsqueeze(1).view(-1) |