from local

This commit is contained in:
dohe0342 2023-01-04 00:12:26 +09:00
parent 9c7d1256f8
commit eff87e2fb6
2 changed files with 7 additions and 1 deletions

View File

@ -188,7 +188,7 @@ class Conformer(EncoderInterface):
warmup=warmup, warmup=warmup,
) # (T, N, C) ) # (T, N, C)
else: else:
x = self.encoder( x, layer_output = self.encoder(
x, x,
pos_emb, pos_emb,
src_key_padding_mask=src_key_padding_mask, src_key_padding_mask=src_key_padding_mask,
@ -197,6 +197,8 @@ class Conformer(EncoderInterface):
x = x.permute(1, 0, 2) # (T, N, C) ->(N, T, C) x = x.permute(1, 0, 2) # (T, N, C) ->(N, T, C)
layer_output = [x.permute(1, 0, 2) for x in layer_output]
return x, lengths return x, lengths
@torch.jit.export @torch.jit.export
@ -693,6 +695,8 @@ class ConformerEncoder(nn.Module):
outputs = [] outputs = []
layer_output = []
for i, mod in enumerate(self.layers): for i, mod in enumerate(self.layers):
output = mod( output = mod(
output, output,
@ -704,6 +708,8 @@ class ConformerEncoder(nn.Module):
if i in self.aux_layers: if i in self.aux_layers:
outputs.append(output) outputs.append(output)
layer_output.append(output)
output = self.combiner(outputs) output = self.combiner(outputs)
return output, layer_output return output, layer_output