added code

code/models/transformer.py

@@ -10,7 +10,8 @@ class TransformerModel(nn.Module):
                  seq_len: int,
                  embed_dim: int,
                  num_heads: int,
-                 num_enc_layers: int):
+                 num_enc_layers: int,
+                 **kwargs):
         super().__init__()
         self.input_proj = nn.Linear(input_dim, embed_dim)
 
@@ -18,7 +19,7 @@ class TransformerModel(nn.Module):
         pe = self._get_sinusoidal_embedding(seq_len, embed_dim)  # (seq_len, embed_dim)
         self.register_buffer('pos_embed', pe.unsqueeze(0))  # (1, seq_len, embed_dim)
 
-        encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads)
+        encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads, dropout=0.1, **kwargs)
         self.encoder = nn.TransformerEncoder(encoder_layer, num_enc_layers)
         self.pool = nn.AdaptiveAvgPool1d(1)
         self.head = nn.Linear(embed_dim, output_dim)