Update ultravox/model/ultravox_model.py

Co-authored-by: Farzad Abdolhosseini <[email protected]>
fixie-ai · Nov 20, 2024 · 93164a8 · 93164a8
1 parent feb2394
commit 93164a8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ultravox/model/ultravox_model.py b/ultravox/model/ultravox_model.py
@@ -544,7 +544,7 @@ class ModifiedWhisperEncoder(
     base_model_prefix = "model.encoder"
     _no_split_modules = ["WhisperEncoderLayer"]
 
-    def init_latency_mask(self, audio_latency_block_size: int, dtype: torch.float):
+    def init_latency_mask(self, audio_latency_block_size: int, dtype: torch.dtype):
         if audio_latency_block_size is None:
             self.audio_streaming_mask = None
             return