update T4 env

app.py

@@ -47,6 +47,7 @@ class TimeSeriesEditor:
             (200, 250, 0, self.function_parser.string_to_function("sin(2*pi*x)"), 0.05)
             # 200,250,0,sin(2*pi*x),0.05
         ]
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     def format_value(self, value: float, feature_idx: int) -> str:
         """Format value with appropriate units and notation"""
@@ -440,6 +441,10 @@ class TimeSeriesEditor:
 
         # Run prediction
         with torch.no_grad():
+            # to cuda
+            observed_points = observed_points.to(self.device)
+            observed_mask = observed_mask.to(self.device)
+
             sample = self.trainer.predict_weighted_points(
                 observed_points, # (seq_length, feature_dim)
                 observed_mask, # (seq_length, feature_dim)
@@ -614,7 +619,7 @@ class TimeSeriesEditor:
 def create_gradio_interface(editor: TimeSeriesEditor):
     with gr.Blocks() as app:
         gr.Markdown("# Time Series Editor")
-        gr.Markdown("## Instruction: Scroll Down + Click [Update Figure] [~20s-30s] [Running on CPU...]")
+        gr.Markdown("## Instruction: Scroll Down + Click [Update Figure] [~20s]")
 
         metrics_display = gr.JSON(label="Metrics", value={})
 
@@ -1073,7 +1078,9 @@ if __name__ == "__main__":
     print(os.getcwd())
 
     device = torch.device(f"cuda:0") if torch.cuda.is_available() else "cpu"
-    
+    print(f"Device: {device}")
+    print(f"Using device: {device}")
+
     from models.Tiffusion import tiffusion
 
     model = tiffusion.Tiffusion(