From 2502987dc48668355d440240dc907316fe555a51 Mon Sep 17 00:00:00 2001 From: yifanyeung Date: Thu, 4 Jul 2024 16:09:20 +0800 Subject: [PATCH] support for cascades --- egs/librispeech/ASR/zipformer_lstm/train.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/egs/librispeech/ASR/zipformer_lstm/train.py b/egs/librispeech/ASR/zipformer_lstm/train.py index 8b3a4644a..f703294a0 100755 --- a/egs/librispeech/ASR/zipformer_lstm/train.py +++ b/egs/librispeech/ASR/zipformer_lstm/train.py @@ -1262,12 +1262,19 @@ def run(rank, world_size, args): sampler_state_dict = None train_dl = librispeech.train_dataloaders( - train_cuts, sampler_state_dict=sampler_state_dict + train_cuts, + sampler_state_dict=sampler_state_dict, + world_size=world_size, + rank=rank, ) valid_cuts = librispeech.dev_clean_cuts() valid_cuts += librispeech.dev_other_cuts() - valid_dl = librispeech.valid_dataloaders(valid_cuts) + valid_dl = librispeech.valid_dataloaders( + valid_cuts, + world_size=world_size, + rank=rank, + ) if 0 and not params.print_diagnostics: scan_pessimistic_batches_for_oom(