From 6cbc02945017c4dc9e8e09a198e9f65d2c537837 Mon Sep 17 00:00:00 2001
From: crackerben <me@crackerben.com>
Date: Wed, 6 Dec 2023 20:34:30 +0800
Subject: [PATCH] Fix timestamp extraction bug in transcribe function

---
 whisper/whisper/transcribe.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/whisper/whisper/transcribe.py b/whisper/whisper/transcribe.py
index ef102a2f..bfdc32b5 100644
--- a/whisper/whisper/transcribe.py
+++ b/whisper/whisper/transcribe.py
@@ -300,7 +300,7 @@ def transcribe(
                     seek += last_timestamp_pos * input_stride
             else:
                 duration = segment_duration
-                timestamps = tokens[timestamp_tokens.nonzero().flatten()]
+                timestamps = tokens[timestamp_tokens.nonzero()[0]]
                 if (
                     len(timestamps) > 0
                     and timestamps[-1].item() != tokenizer.timestamp_begin