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

Merge branch 'scaled_adam_exp151' into scaled_adam_exp150

Browse Source
This commit is contained in:
Daniel Povey 2022-10-19 19:59:07 +08:00
commit f08a869769
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

@ -255,10 +255,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):