further cleanup and resetting device to cuda inside the process function

app.py

@@ -18,11 +18,7 @@ model_card = ModelCard(
 
 def load_audio(file_path):
     audio, samplerate = torchaudio.load(file_path)
-    return audio.unsqueeze(0)  # [1, 1, samples] - no .cuda()
-
-def save_audio(file_path, audio, samplerate=44100):
-    audio = audio.squeeze(0).cpu()
-    torchaudio.save(file_path, audio, samplerate)
+    return audio.unsqueeze(0) 
 
 # Load the model outside of the process function so that it only has to happen once
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -76,6 +72,7 @@ def process_fn(
     input_audio_path: str
 ) -> str:
     
+    device = torch.device("cuda")
     sig = load_audio(input_audio_path)
 
     # Move audio data to device
@@ -85,8 +82,7 @@ def process_fn(
     if sig.dim() == 2:
         sig = sig.unsqueeze(0)
     
-    with torch.inference_mode():
-        result = model(sig)
+    result = model(sig)
 
     # Remove batch dimension
     result = result.squeeze(0)