.

app.py

@@ -4,6 +4,7 @@ from train import train as train_fn
 from testing import evaluate as eval_fn
 from inference import predict_next as predict_fn
 
+@spaces.GPU
 def train_api(symbol, seq_len=60, epochs=5, batch_size=32, start="", end=""):
     return train_fn(
         symbol,

inference.py

@@ -8,6 +8,7 @@ from models import StockLSTM
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 ARTIFACTS_DIR = "artifacts"
+device = torch.device("cpu")
 
 def _paths(symbol: str):
     base = os.path.join(ARTIFACTS_DIR, symbol.upper())
@@ -27,7 +28,8 @@ def _load_artifacts(symbol: str):
         scaler = pickle.load(f)
 
     model = StockLSTM(input_dim=1, hidden_dim=64, num_layers=2, dropout=0.2)
-    model.load_state_dict(torch.load(p["model"], map_location="cpu"))
+    model = model.to(device)
+    model.load_state_dict(torch.load(p["model"], map_location="cpu"))    
     model.eval()
     return model, scaler, meta