Audio Encoder to bfloat16 (#4)

* audio enc to bfloat16 * increase transformers min required version
fixie-ai · Jun 3, 2024 · 56a8209 · 56a8209
1 parent d2df97b
commit 56a8209
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 # Core
-transformers[torch]>=4.39.3
+transformers[torch]>=4.40.2
 bitsandbytes>=0.42.0
 peft
 simple_parsing

diff --git a/ultravox/training/train.py b/ultravox/training/train.py
@@ -174,9 +174,7 @@ def main() -> None:
     logging.info(
         f"Using dtype and device (world_size): {dtype}, {device} ({world_size})"
     )
-    model.to(device)
-    model.language_model.to(dtype)
-    model.multi_modal_projector.to(dtype)
+    model.to(device=device, dtype=dtype)
     # TODO: check if the whole model can now be moved to dtype instead
 
     # Prepare dataset, subsetting if needed