This commit is contained in:
嘉渊 2023-07-29 15:13:02 +08:00
parent 6835108df0
commit 670c9f3d39

View File

@ -12,12 +12,12 @@ from typing import Tuple
import torch import torch
from funasr.layers.abs_normalize import AbsNormalize from funasr.layers.abs_normalize import AbsNormalize
from funasr.models.base_model import FunASRModel
from funasr.models.encoder.abs_encoder import AbsEncoder from funasr.models.encoder.abs_encoder import AbsEncoder
from funasr.models.frontend.abs_frontend import AbsFrontend from funasr.models.frontend.abs_frontend import AbsFrontend
from funasr.models.preencoder.abs_preencoder import AbsPreEncoder from funasr.models.preencoder.abs_preencoder import AbsPreEncoder
from funasr.models.specaug.abs_specaug import AbsSpecAug from funasr.models.specaug.abs_specaug import AbsSpecAug
from funasr.torch_utils.device_funcs import force_gatherable from funasr.torch_utils.device_funcs import force_gatherable
from funasr.models.base_model import FunASRModel
if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
from torch.cuda.amp import autocast from torch.cuda.amp import autocast
@ -36,8 +36,8 @@ class Data2VecPretrainModel(FunASRModel):
frontend: Optional[AbsFrontend], frontend: Optional[AbsFrontend],
specaug: Optional[AbsSpecAug], specaug: Optional[AbsSpecAug],
normalize: Optional[AbsNormalize], normalize: Optional[AbsNormalize],
preencoder: Optional[AbsPreEncoder],
encoder: AbsEncoder, encoder: AbsEncoder,
preencoder: Optional[AbsPreEncoder] = None,
): ):
super().__init__() super().__init__()