summaryrefslogtreecommitdiff
path: root/models/auto_encoder.py
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2020-12-23 20:15:14 +0800
committerJordan Gong <jordan.gong@protonmail.com>2020-12-23 20:15:14 +0800
commit7460ce2c8904a009f2f1139b11ec18faf208d6d2 (patch)
tree6f394b4072ba2c670c1eaabd4da1f2a1d0c8b047 /models/auto_encoder.py
parentfa745aa38edce88f5d92e7f7242a111c3d892161 (diff)
Reshape feature before decode
Diffstat (limited to 'models/auto_encoder.py')
-rw-r--r--models/auto_encoder.py3
1 files changed, 2 insertions, 1 deletions
diff --git a/models/auto_encoder.py b/models/auto_encoder.py
index e35ed23..1708bc9 100644
--- a/models/auto_encoder.py
+++ b/models/auto_encoder.py
@@ -79,9 +79,10 @@ class Decoder(nn.Module):
def forward(self, fa, fgs, fgd):
x = torch.cat((fa, fgs, fgd), dim=1).view(-1, self.em_dim)
x = F.leaky_relu(self.batch_norm_fc(self.fc(x)), 0.2)
+ x = x.view(-1, 64 * 8, 4, 2)
x = F.leaky_relu(self.batch_norm1(self.trans_conv1(x)), 0.2)
x = F.leaky_relu(self.batch_norm2(self.trans_conv2(x)), 0.2)
x = F.leaky_relu(self.batch_norm3(self.trans_conv3(x)), 0.2)
x = F.sigmoid(self.trans_conv4(x))
- return x \ No newline at end of file
+ return x