fixed hard-coded device

speechbrain · Aug 14, 2023 · 091b3ce · 091b3ce
1 parent 331acdb
commit 091b3ce
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/recipes/SLURP/direct/hparams/train.yaml b/recipes/SLURP/direct/hparams/train.yaml
@@ -60,9 +60,7 @@ epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter
     limit: !ref <number_of_epochs>
 
 # Models
-asr_model: !apply:speechbrain.pretrained.EncoderDecoderASR.from_hparams
-    source: speechbrain/asr-crdnn-rnnlm-librispeech
-    run_opts: {"device":"cuda:0"}
+asr_model_source: speechbrain/asr-crdnn-rnnlm-librispeech
 
 slu_enc: !new:speechbrain.nnet.containers.Sequential
     input_shape: [null, null, !ref <ASR_encoder_dim>]

diff --git a/recipes/SLURP/direct/train.py b/recipes/SLURP/direct/train.py
@@ -340,6 +340,14 @@ def text_pipeline(semantics):
     run_on_main(hparams["pretrainer"].collect_files)
     hparams["pretrainer"].load_collected(device=run_opts["device"])
 
+    # Download pretrained ASR model
+    from speechbrain.pretrained import EncoderDecoderASR
+
+    hparams["asr_model"] = EncoderDecoderASR.from_hparams(
+        source=hparams["asr_model_source"],
+        run_opts={"device": run_opts["device"]},
+    )
+
     # Brain class initialization
     slu_brain = SLU(
         modules=hparams["modules"],