diff --git a/whisperlivekit/diarization/sortformer_backend.py b/whisperlivekit/diarization/sortformer_backend.py
index dd35cfd..3072eae 100644
--- a/whisperlivekit/diarization/sortformer_backend.py
+++ b/whisperlivekit/diarization/sortformer_backend.py
@@ -106,6 +106,7 @@ class SortformerDiarizationOnline:
             features=128,
             pad_to=0
         )
+        self.audio2mel.to(self.diar_model.device)
         
         self.chunk_duration_seconds = (
             self.diar_model.sortformer_modules.chunk_len *