Update app.py

app.py

@@ -128,14 +128,15 @@ class GRUModel(nn.Module):
 class BiLSTMModel(nn.Module):
     def __init__(self):
         super(BiLSTMModel, self).__init__()
-        self.lstm = nn.LSTM(input_size=1, hidden_size=100, num_layers=2, batch_first=True, dropout=0.2, bidirectional=True)
+        self.lstm = nn.LSTM(input_size=1, hidden_size=100, num_layers=1, batch_first=True, dropout=0.2, bidirectional=True)
         self.fc = nn.Linear(200, 1)
 
     def forward(self, x):
-        h0 = torch.zeros(4, x.size(0), 100)  # 2 directions × 2 layers = 4
-        c0 = torch.zeros(4, x.size(0), 100)
+        h0 = torch.zeros(2, x.size(0), 100)
+        c0 = torch.zeros(2, x.size(0), 100)
         out, _ = self.lstm(x, (h0, c0))
         return self.fc(out[:, -1, :])
+
 @st.cache_resource(ttl=3600)
 def load_model_from_disk(path, model_type):
     model = BiLSTMModel() if model_type == "Bi-Directional LSTM" else GRUModel()