雾聪
2023-06-28 54931dd4e1a099d7d6f144c4e12e5453deb3aa26
funasr/models/e2e_diar_eend_ola.py
@@ -16,7 +16,7 @@
from funasr.modules.eend_ola.encoder_decoder_attractor import EncoderDecoderAttractor
from funasr.modules.eend_ola.utils.power import generate_mapping_dict
from funasr.torch_utils.device_funcs import force_gatherable
from funasr.train.abs_espnet_model import AbsESPnetModel
from funasr.models.base_model import FunASRModel
if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
    pass
@@ -34,7 +34,7 @@
    return att
class DiarEENDOLAModel(AbsESPnetModel):
class DiarEENDOLAModel(FunASRModel):
    """EEND-OLA diarization model"""
    def __init__(
@@ -91,7 +91,6 @@
            text_lengths: torch.Tensor,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
        """Frontend + Encoder + Decoder + Calc loss
        Args:
            speech: (Batch, Length, ...)
            speech_lengths: (Batch, )