More recurrence fixes for chainedgen

2020-10-17 08:35:46 -06:00 · 2020-10-17 08:35:46 -06:00 · 6141aa1110
commit 6141aa1110
parent cf8118a85b
1 changed files with 2 additions and 1 deletions
--- a/codes/models/archs/ChainedEmbeddingGen.py
+++ b/codes/models/archs/ChainedEmbeddingGen.py
@ -81,9 +81,10 @@ class ChainedEmbeddingGenWithStructure(nn.Module):
        self.grad_extract = ImageGradientNoPadding()
        self.upsample = FinalUpsampleBlock2x(64)

-    def forward(self, x):
+    def forward(self, x, recurrent=None):
        emb = checkpoint(self.spine, x)
        if self.recurrent:
+            fea = torch.cat([x,recurrent], dim=1)
            fea = self.initial_conv_rec(x)
        else:
            fea = self.initial_conv(x)