small changes

arbit3rr · Oct 23, 2023 · e239c31 · e239c31
1 parent 2035312
commit e239c31
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/config.yaml b/config.yaml
@@ -3,9 +3,10 @@
 # Dataset settings
 NUM_CYCLES: 3
 FEATURE_DIM: 16
+NUM_ATTENTION: 8
 
 # Training settings
-EPOCHS: 1000
+EPOCHS: 2000
 LEARNING_RATE: 0.001
 
 # Other settings

diff --git a/model.py b/model.py
@@ -9,7 +9,7 @@ def __init__(self, feature_dim, num_cycles):
         self.final_linear_layer = nn.Linear(feature_dim, 1)
 
         # self.conv_layer = nn.Conv1d(3, 512, kernel_size=16, stride=8)
-        self.conv_layer = nn.Conv2d(num_cycles, feature_dim, kernel_size=3, stride=8)
+        self.conv_layer = nn.Conv2d(num_cycles, feature_dim, kernel_size=3, stride=4)
         self.encoder_layer = nn.TransformerEncoderLayer(d_model=feature_dim, nhead=8, batch_first=True)
         self.decoder_layer = nn.TransformerDecoderLayer(d_model=feature_dim, nhead=8, batch_first=True)