feat: improve usage of GPU

app.py

@@ -7,6 +7,7 @@ import numpy as np
 import gradio as gr
 import asyncio
 import spaces
+import torch
 
 # Set matplotlib to use non-interactive backend before importing pyplot
 # This is required for Gradio which runs on worker threads
@@ -55,7 +56,7 @@ def load_pipeline(model_id: str, species: str = DEFAULT_SPECIES):
     pipe = load_ntv3_tracks_pipeline(
         model=model_id,
         token=HF_TOKEN,
-        device="auto",
+        device="cpu",  # This prevents the pipeline constructor from doing model.to(cuda) during import.
         default_species=species,
         verbose=False,
     )
@@ -427,8 +428,19 @@ def predict(
     if "species" not in inputs:
         raise gr.Error(f"Internal error: species not found in inputs dict. Inputs: {list(inputs.keys())}")
 
+    # move to GPU only once the ZeroGPU context is active
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe.model.to(device)
+    pipe.model.eval()
+    print(f"Running on {next(pipe.model.parameters()).device}")
+
+    # run inference
     out = pipe(inputs)
 
+    # optional: move back to CPU so you don’t rely on any persistent CUDA context
+    if device == "cuda":
+        pipe.model.to("cpu")
+
     bw_names = out.bigwig_track_names or []
     bw = out.bigwig_tracks_logits
     bed_names = out.bed_element_names or []