Update model.py

model.py

@@ -67,16 +67,19 @@ class INF5Model(PreTrainedModel):
         self.ema_model.load_state_dict(state_dict, strict=False)
 
     def _extract_embedding_from_audio_and_text(self, audio_path: str, text: str) -> torch.Tensor:
-    
         device = next(self.parameters()).device  # model device
     
-        # Load audio waveform
+        # Load audio waveform on CPU first
         waveform, sample_rate = torchaudio.load(audio_path)
         target_sample_rate = 24000
         if sample_rate != target_sample_rate:
+            # Move waveform to device before resampling to avoid device mismatch
+            waveform = waveform.to(device)
             resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate).to(device)
             waveform = resampler(waveform)
-        waveform = waveform.to(device)
+        else:
+            # If no resampling, still move waveform to device for model
+            waveform = waveform.to(device)
     
         # Forward pass - pass waveform and text directly to ema_model
         with torch.no_grad():