diff --git a/modules/diarize/diarize_pipeline.py b/modules/diarize/diarize_pipeline.py
index bdb7d53..4313f5c 100644
--- a/modules/diarize/diarize_pipeline.py
+++ b/modules/diarize/diarize_pipeline.py
@@ -44,9 +44,9 @@ def __call__(self, audio: Union[str, np.ndarray], min_speakers=None, max_speaker
 
 def assign_word_speakers(diarize_df, transcript_result, fill_nearest=False):
     transcript_segments = transcript_result["segments"]
+    if transcript_segments and isinstance(transcript_segments[0], Segment):
+        transcript_segments = [seg.model_dump() for seg in transcript_segments]
     for seg in transcript_segments:
-        if isinstance(seg, Segment):
-            seg = seg.model_dump()
         # assign speaker to segment (if any)
         diarize_df['intersection'] = np.minimum(diarize_df['end'], seg['end']) - np.maximum(diarize_df['start'],
                                                                                             seg['start'])
@@ -88,7 +88,7 @@ def assign_word_speakers(diarize_df, transcript_result, fill_nearest=False):
                     if word_speaker is not None:
                         word["speaker"] = word_speaker
 
-    return transcript_result
+    return {"segments": transcript_segments}
 
 
 class DiarizationSegment:
diff --git a/modules/diarize/diarizer.py b/modules/diarize/diarizer.py
index 4c0727a..38e150a 100644
--- a/modules/diarize/diarizer.py
+++ b/modules/diarize/diarizer.py
@@ -71,7 +71,6 @@ def run(self,
 
         segments_result = []
         for segment in diarized_result["segments"]:
-            segment = segment.dict()
             speaker = "None"
             if "speaker" in segment:
                 speaker = segment["speaker"]