debug

alexandrainst · Sep 27, 2024 · b7a4ed1 · b7a4ed1
1 parent 3088406
commit b7a4ed1
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/coral/finetune.py b/src/coral/finetune.py
@@ -5,10 +5,10 @@
 
 from omegaconf import DictConfig
 from transformers import (
-    AutoModelForCTC,
-    AutoProcessor,
     EarlyStoppingCallback,
     TrainerCallback,
+    Wav2Vec2ForCTC,
+    Wav2Vec2ProcessorWithLM,
 )
 from wandb.sdk.wandb_init import init as wandb_init
 
@@ -37,8 +37,8 @@ def finetune(config: DictConfig) -> None:
     dataset = load_data_for_finetuning(config=config, processor=processor)
 
     # TEMP
-    model = AutoModelForCTC.from_pretrained(config.model_dir)
-    processor = AutoProcessor.from_pretrained(config.model_dir)
+    model = Wav2Vec2ForCTC.from_pretrained(config.model_dir)
+    processor = Wav2Vec2ProcessorWithLM.from_pretrained(config.model_dir)
     breakpoint()
 
     if config.wandb and is_main_process: