diff --git a/README.md b/README.md
index ff9d2aa..e62426b 100644
--- a/README.md
+++ b/README.md
@@ -166,6 +166,7 @@ See more examples in other languages [here](EXAMPLES.md).
 ```python
 import whisperx
 import gc
+from whisperx.diarize import DiarizationPipeline
 
 device = "cuda"
 audio_file = "audio.mp3"
@@ -196,7 +197,7 @@ print(result["segments"]) # after alignment
 # import gc; import torch; gc.collect(); torch.cuda.empty_cache(); del model_a
 
 # 3. Assign speaker labels
-diarize_model = whisperx.diarize.DiarizationPipeline(use_auth_token=YOUR_HF_TOKEN, device=device)
+diarize_model = DiarizationPipeline(use_auth_token=YOUR_HF_TOKEN, device=device)
 
 # add min/max number of speakers if known
 diarize_segments = diarize_model(audio)