mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-12-11 06:55:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use full precision to do softmax and store ans.

Browse Source
This commit is contained in:
Daniel Povey 2022-10-19 19:53:53 +08:00
parent 0ad4462632
commit cc15552510
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py
View File

@ -219,10 +219,13 @@ class SoftmaxFunction(torch.autograd.Function):
"""
@staticmethod
def forward(ctx, x: Tensor, dim: int):
ans = x.softmax(dim=dim)
ctx.save_for_backward(ans)
ctx.dim = dim
return ans
with torch.cuda.amp.autocast(enabled=False):
if x.dtype == torch.float16:
x = x.to(torch.float32)
ans = x.softmax(dim=dim)
ctx.save_for_backward(ans)
ctx.dim = dim
return ans
@staticmethod
def backward(ctx, ans_grad: Tensor):