working static ae training

2023-04-06 17:25:57 +00:00
parent 31c2194a4e
commit b4be8bf80d
5 changed files with 172 additions and 25 deletions
@@ -69,7 +69,8 @@ class Decoder(nn.Module):

        self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_classes, device=device, reversed=True)
            
-        self.softmax = nn.Softmax(dim=1)
+        if self.is_label:
+            self.softmax = nn.Softmax(dim=1)

    def forward(self, inputs):
        if len(inputs) == 2: