Update models/av_mossformer2_tse/av_mossformer2.py

models/av_mossformer2_tse/av_mossformer2.py

@@ -200,9 +200,9 @@ class AV_MossFormer2_TSE_16K(nn.Module):
 
     def __init__(self, args):
         super(AV_MossFormer2_TSE_16K, self).__init__()
-        # self.model = av_mossformer2(args)
-        from .av_mossformer_tmp import av_mossformer_tmp
-        self.model = av_mossformer_tmp(args)
+        self.model = av_mossformer2(args)
+        # from .av_mossformer_tmp import av_mossformer_tmp
+        # self.model = av_mossformer_tmp(args)
 
     def forward(self, x):
         outputs = self.model(x)