Why cpu and gpu have different results when logit_length and target_length has some 1？

🐛 Describe the bug

Some elements of logit_length and target_length are 1 respectively，The calculation results of cpu and gpu will be different，Is gpu‘s wrong or is call method of my wrong?

test_code:

import torch
import torchaudio
import numpy as np 
np.random.seed(0)
rnnt_loss_betas = torch._C._jit_get_operation('torchaudio::rnnt_loss_betas')[0]
rnnt_loss = torch._C._jit_get_operation('torchaudio::rnnt_loss')[0]
B,T,U,D = 4,2,3,5
logits_cpu = torch.tensor(np.random.rand(B,T,U,D),dtype = torch.float32)

target_cpu = torch.tensor(np.random.randint(1,D,[B,U-1]),dtype=torch.int32)

logit_length_cpu = torch.tensor([2,1,1,1],dtype=torch.int32)
target_length_cpu = torch.tensor([2,1,1,1],dtype=torch.int32)

beta_cpu = rnnt_loss_betas(logits_cpu,target_cpu,logit_length_cpu,target_length_cpu,0,-1)
cost_cpu,grad_cpu = rnnt_loss(logits_cpu,target_cpu,logit_length_cpu,target_length_cpu,0,1)

logits = logits_cpu.cuda()
targets = target_cpu.cuda()
logit_length = logit_length_cpu.cuda()
target_length = target_length_cpu.cuda()

beta_gpu = rnnt_loss_betas(logits,targets,logit_length,target_length,0,-1)
cost_gpu,grad_gpu = rnnt_loss(logits,targets,logit_length,target_length,0,-1)

print((logits_cpu == logits.cpu()).all())
print((target_cpu == targets.cpu()).all())
print((logit_length_cpu == logit_length.cpu()).all())
# print("cpu = {} \ngpu = {}".format(beta_cpu,beta_gpu))
print((abs(beta_cpu-beta_gpu.cpu())<0.001).all())
print((abs(cost_cpu-cost_gpu.cpu())<0.001).all())
print((abs(grad_cpu-grad_gpu.cpu())<0.001).all())

output：

tensor(True)
tensor(True)
tensor(True)
tensor(False)
tensor(False)
tensor(False)

Versions

https://gist.github.com/bleedingfight/f31ffe7dc30813151ef9664077a331f5