Fix type hint

This commit is contained in:
yfy62 2023-04-26 17:43:11 +08:00
parent 0b5996bd3d
commit 23a9b66295

View File

@ -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,