diff --git a/whisperx/diarize.py b/whisperx/diarize.py
index d59a847..55fe220 100644
--- a/whisperx/diarize.py
+++ b/whisperx/diarize.py
@@ -26,10 +26,9 @@ class DiarizationPipeline:
             'sample_rate': SAMPLE_RATE
         }
         segments = self.model(audio_data, min_speakers=min_speakers, max_speakers=max_speakers)
-        diarize_df = pd.DataFrame(segments.itertracks(yield_label=True))
-        diarize_df['start'] = diarize_df[0].apply(lambda x: x.start)
-        diarize_df['end'] = diarize_df[0].apply(lambda x: x.end)
-        diarize_df.rename(columns={2: "speaker"}, inplace=True)
+        diarize_df = pd.DataFrame(segments.itertracks(yield_label=True), columns=['segment', 'label', 'speaker'])
+        diarize_df['start'] = diarize_df['segment'].apply(lambda x: x.start)
+        diarize_df['end'] = diarize_df['segment'].apply(lambda x: x.end)
         return diarize_df