Update nets/decoder.py

nets/decoder.py

@@ -79,10 +79,10 @@ class Attention(nn.Module):
 
 
 class Decoder(nn.Module):
-    def __init__(self, num_heads, input_size, embedding_size, softmax_output=False, C=10):
+    def __init__(self, num_heads, embedding_size, decoder_input_size, softmax_output=False, C=10):
         super().__init__()
         self.embedding_size = embedding_size
-        self.initial_embedding = nn.Linear(input_size, embedding_size)
+        self.initial_embedding = nn.Linear(decoder_input_size, embedding_size)
         self.attention = Attention(n_heads=num_heads, input_dim=embedding_size, embed_dim=embedding_size)
         self.classifier_output = ClassifierOutput(embedding_size=embedding_size, C=C, softmax_output=softmax_output)