Skip to content

Commit

Permalink
Fix exporting emformer with torchscript using torch 1.6.0 (k2-fsa#402)
Browse files Browse the repository at this point in the history
  • Loading branch information
csukuangfj authored Jun 7, 2022
1 parent 29fa878 commit 80c46f0
Showing 1 changed file with 3 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,7 @@ def __init__(
)
self.log_eps = math.log(1e-10)

self._has_init_state = False
self._init_state = torch.jit.Attribute([], List[List[torch.Tensor]])

def forward(
Expand Down Expand Up @@ -296,7 +297,7 @@ def get_init_state(self, device: torch.device) -> List[List[torch.Tensor]]:
Return the initial state of each layer. NOTE: the returned
tensors are on the given device. `len(ans) == num_emformer_layers`.
"""
if len(self._init_state) > 0:
if self._has_init_state:
# Note(fangjun): It is OK to share the init state as it is
# not going to be modified by the model
return self._init_state
Expand All @@ -308,6 +309,7 @@ def get_init_state(self, device: torch.device) -> List[List[torch.Tensor]]:
s = layer._init_state(batch_size=batch_size, device=device)
ans.append(s)

self._has_init_state = True
self._init_state = ans

return ans

0 comments on commit 80c46f0

Please sign in to comment.