From 23a9b662950c6ead010ea45ba07515a4d30857d4 Mon Sep 17 00:00:00 2001
From: yfy62 <yfy62@d3-hpc-sjtu-test-003.cm.cluster>
Date: Wed, 26 Apr 2023 17:43:11 +0800
Subject: [PATCH] Fix type hint

---
 egs/gigaspeech/ASR/pruned_transducer_stateless2/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/egs/gigaspeech/ASR/pruned_transducer_stateless2/train.py b/egs/gigaspeech/ASR/pruned_transducer_stateless2/train.py
index 6f9a85c2d..32998a8de 100755
--- a/egs/gigaspeech/ASR/pruned_transducer_stateless2/train.py
+++ b/egs/gigaspeech/ASR/pruned_transducer_stateless2/train.py
@@ -678,7 +678,7 @@ def compute_validation_loss(
 
 def train_one_epoch(
     params: AttributeDict,
-    model: nn.Module,
+    model: Union[nn.Module, DDP],
     optimizer: torch.optim.Optimizer,
     scheduler: LRSchedulerType,
     sp: spm.SentencePieceProcessor,