.

.gitignore

@@ -1 +1,2 @@
 __pycache__/
+artifacts/

inference.py

@@ -6,6 +6,7 @@ import torch
 
 from models import StockLSTM
 
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
 ARTIFACTS_DIR = "artifacts"
 
 def _paths(symbol: str):

models.py

@@ -1,6 +1,9 @@
+import os
 import torch
 import torch.nn as nn
 
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
 class StockLSTM(nn.Module):
     def __init__(self, input_dim=1, hidden_dim=64, num_layers=2, dropout=0.2):
         super().__init__()

testing.py

@@ -7,6 +7,7 @@ import torch
 
 from models import StockLSTM
 
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
 ARTIFACTS_DIR = "artifacts"
 
 @torch.no_grad()