Upload core/train_eval.py with huggingface_hub

core/train_eval.py

@@ -1,4 +1,3 @@
-import spaces
 import numpy as np
 import pandas as pd
 import torch
@@ -18,6 +17,7 @@ import logging
 import torch.optim.lr_scheduler as lr_scheduler
 from io import StringIO
 import sys
+import spaces
 
 try:
     from torchsummary import summary
@@ -123,7 +123,7 @@ def select_features(df, features, target, selector_method, importance_threshold)
         logging.warning(f"Unsupported selector_method: {selector_method}, using all features")
         return features
 
-@spaces.GPU(duration=60)  # Adjust duration as needed
+@spaces.GPU
 def train_and_evaluate(
     df,
     features,
@@ -143,13 +143,10 @@ def train_and_evaluate(
     selector_method="RandomForest",
     importance_threshold=0.0,
     scheduler_type="None",
-    device=None,
+    device='cuda',
     verbose=True
 ):
     try:
-        # Ensure device is set to cuda if available
-        if device is None:
-            device = 'cuda' if torch.cuda.is_available() else 'cpu'
         logging.info(f"Starting train_and_evaluate: model={model_cls.__name__}, features={len(features)}, window={window}, horizon={horizon}, scheduler={scheduler_type}, selector_method={selector_method}")
         from .data import preprocess_data