.

train.py

@@ -44,7 +44,7 @@ def to_tensor_loader(X, y, batch_size=32):
     return ds
 
 def train(symbol: str, seq_len: int = 60, epochs: int = 5, batch_size: int = 32,
-          start: str = None, end: str = None, lr: float = 1e-3):
+          start: str = None, end: str = None, lr: float = 1e-2):
     device = torch.device("cpu")
 
     # --- data ---
@@ -75,8 +75,8 @@ def train(symbol: str, seq_len: int = 60, epochs: int = 5, batch_size: int = 32,
     # --- model ---
     model = StockLSTM(input_dim=1, hidden_dim=64, num_layers=2, dropout=0.2).to(device)
     criterion = nn.MSELoss()
-    # optimizer = optim.Adam(model.parameters(), lr=lr)
-    optimizer = optim.SGD(model.parameters(), lr=lr)
+    optimizer = optim.AdamW(model.parameters(), lr=lr)
+    # optimizer = optim.SGD(model.parameters(), lr=lr)
 
     # --- training ---
     model.train()