Forgot to add audio to current device.

app.py

@@ -22,8 +22,8 @@ model_path_orig = os.path.join('models', 'tt-orig.pt')
 tt_weights_orig = torch.load(model_path_orig, map_location='cpu')
 #tt_weights_demo = torch.load(model_path_demo, map_location='cpu')
 
-if torch.cuda.is_available():
-    model = model.cuda()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device)
 
 model_card = ModelCard(
     name='Timbre-Trap',
@@ -32,7 +32,6 @@ model_card = ModelCard(
     tags=['example', 'music transcription', 'multi-pitch estimation', 'timbre filtering']
 )
 
-
 def process_fn(audio_path, transcribe):#, demo):
     # Load the audio with torchaudio
     audio, fs = torchaudio.load(audio_path)
@@ -54,6 +53,9 @@ def process_fn(audio_path, transcribe):#, demo):
     # Load weights of the original model
     model.load_state_dict(tt_weights_orig)
 
+    # Add audio to current device
+    audio = audio.to(device)
+
     # Obtain transcription or reconstructed spectral coefficients
     coefficients = model.chunked_inference(audio, transcribe)
 
@@ -70,6 +72,9 @@ def process_fn(audio_path, transcribe):#, demo):
     # Resample audio back to the original sampling rate
     audio = torchaudio.functional.resample(audio, 22050, fs)
 
+    # Bring audio back to CPU
+    audio = audio.cpu()
+
     # Create a temporary directory for output
     os.makedirs('_outputs', exist_ok=True)
     # Create a path for saving the audio