move waveform to device

app.py

@@ -21,8 +21,9 @@ STEM_CHOICES = {
 
 @spaces.GPU
 def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> AudioSignal:
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
     model = pretrained.get_model(model_name)
-    model.to('cuda' if torch.cuda.is_available() else 'cpu')
+    model.to(device)
     model.eval()
 
     waveform, sr = torchaudio.load(audio_file_path)