added code

code/models/lstm.py

@@ -1,29 +1,22 @@
+import torch
 from torch import nn
 
 
-class LSTMModel(nn.Module):
+class LSTM(nn.Module):
     def __init__(self,
                  input_dim: int,
                  output_dim: int,
-                 cnn_channels: int,
-                 cnn_kernel_size: int,
-                 embed_dim: int,
                  lstm_hidden_size=64,
-                 num_layers=1):
+                 num_layers=1,
+                 **kwargs):
         super().__init__()
-        self.cnn = nn.Sequential(
-            nn.Conv1d(input_dim, cnn_channels, kernel_size=cnn_kernel_size, padding=1),
-            nn.BatchNorm1d(cnn_channels),
-            nn.ReLU(),
-            nn.Conv1d(cnn_channels, embed_dim, kernel_size=cnn_kernel_size, padding=1),
-            nn.BatchNorm1d(embed_dim),
-            nn.ReLU()
-        )
         self.lstm = nn.LSTM(
-            input_size=embed_dim,
+            input_size=input_dim,
             hidden_size=lstm_hidden_size,
             num_layers=num_layers,
             batch_first=True,
+            dropout=0.5,
+            **kwargs
         )
         self.head = nn.Sequential(
             nn.Linear(lstm_hidden_size, output_dim)  # Output is a scalar Δt
@@ -31,12 +24,6 @@ class LSTMModel(nn.Module):
 
     def forward(self, x):
         # x: (batch_size, seq_len, input_size)
-        x = x.permute(0, 2, 1)
-        # x: (batch_size, input_size, seq_len)
-        x = self.cnn(x)
-        # x: (batch_size, embed_dim, seq_len)
-        x = x.permute(0, 2, 1)
-        # x: (batch_size, seq_len, embed_dim)
         x, _ = self.lstm(x)
         # x: (batch_size, seq_len, lstm_hidden_size)
         x = x[:, -1, :]  # Get the last time step