diff options
| author | Jordan Gong <jordan.gong@protonmail.com> | 2021-03-02 20:22:38 +0800 |
|---|---|---|
| committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-03-02 20:22:38 +0800 |
| commit | b274351a8528bc63e52afd6a5d0c34811cea84b1 (patch) | |
| tree | b41df93b460f2166f930fad5ef34ca92d3dc775f /models/rgb_part_net.py | |
| parent | c3143f388730d2869067f6f259775289c742bb48 (diff) | |
| parent | 7fac206f92602462ad8eecde524b0324f7991bde (diff) | |
Merge branch 'data_parallel' into data_parallel_py3.8
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 310ef25..1cda91c 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -69,7 +69,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) |