remove zerogpu

app.py

@@ -3,7 +3,7 @@ from typing import Optional, Iterable
 
 # Third-party library imports
 from transformers import PretrainedConfig, AutoformerForPrediction
-from functools import lru_cache, partial
+from functools import partial
 
 import gradio as gr
 import spaces
@@ -245,15 +245,8 @@ def plot(ts_index, test_dataset, forecasts, prediction_length):
     plt.legend()
     return fig
 
-
-
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
-
-@spaces.GPU
 def do_prediction(days_to_predict: int):
-    device = zero.device
-
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     # Define the desired prediction length
     prediction_length = 7  # Number of time steps to predict into the future
     freq = "1D"  # Daily frequency