summaryrefslogtreecommitdiff
path: root/models/rgb_part_net.py
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-03-14 21:08:02 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-03-14 21:08:02 +0800
commit17ae5e529475a7e47fdde0ce69982b5dfe34f2eb (patch)
tree2293f924b7af940c3c718cc421732d3270d95fd6 /models/rgb_part_net.py
parent3a69cb5325b4a8a51b75224906d839863adeddd6 (diff)
parentda922be042d96338a3f207386e410b6746d046f5 (diff)
Merge branch 'master' into python3.8
Diffstat (limited to 'models/rgb_part_net.py')
-rw-r--r--models/rgb_part_net.py2
1 files changed, 1 insertions, 1 deletions
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index 6e72e38..fba3485 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -59,7 +59,7 @@ class RGBPartNet(nn.Module):
if self.training:
return x_c, x_p, ae_losses, images
else:
- return torch.cat((x_c, x_p)).unsqueeze(1).view(-1)
+ return torch.cat((x_c.view(-1), x_p.view(-1)))
def _disentangle(self, x_c1_t2, x_c2_t2=None):
n, t, c, h, w = x_c1_t2.size()