Testing CUDA

README.md

@@ -1,5 +1,5 @@
 ---
-title: Apollo-CPU
+title: Apollo
 emoji: 💻
 colorFrom: yellow
 colorTo: yellow

app.py

@@ -25,14 +25,17 @@ def save_audio(file_path, audio, samplerate=44100):
     audio = audio.squeeze(0).cpu()
     torchaudio.save(file_path, audio, samplerate)
 
-#Defining the process function
+# Load the model outside of the process function so that it only has to happen once
+
+
+# Defining the process function
 @spaces.GPU
 @torch.inference_mode()
 def process_fn(
     input_audio_path: str
 ) -> str:
-    # Don't set CUDA device - let it use CPU
-    device = torch.device("cpu")
+    # Set CUDA device; use CPU as backup
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
     print(f"Using device: {device}")
     print("Loading Apollo model...")