mirrors/icefall
1
0
Fork 0
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 @staticmethod
def forward(ctx, x: Tensor, dim: int): def forward(ctx, x: Tensor, dim: int):
ans = x.softmax(dim=dim) with torch.cuda.amp.autocast(enabled=False):
ctx.save_for_backward(ans) if x.dtype == torch.float16:
ctx.dim = dim x = x.to(torch.float32)
return ans ans = x.softmax(dim=dim)
ctx.save_for_backward(ans)
ctx.dim = dim
return ans
@staticmethod @staticmethod
def backward(ctx, ans_grad: Tensor): def backward(ctx, ans_grad: Tensor):