Update nets/model.py

nets/model.py

@@ -7,6 +7,7 @@ from nets.encoder import Encoder
 
 class Model(nn.Module):
     def __init__(self, input_size, embedding_size,
+                 decoder_input_size,
                  num_heads=8, num_layers=4, ff_hidden=250, *args, **kwargs):
         super().__init__()
 
@@ -22,9 +23,8 @@ class Model(nn.Module):
         )
 
         # ----------- Decoder -----------
-        decoder_input_dim = 3 * embedding_size + 2  # example: [encoder output] + [mask info] + [fleet info]
         self.decoder = Decoder(
-            input_size=decoder_input_dim,
+            decoder_input_size=decoder_input_size,
             embedding_size=embedding_size,
             num_heads=num_heads
         )