I have tried some other implementation like the following one,
def own_softmax(self, x)
    
    means = torch.mean(x, 1, keepdim=True)[0]
    x_exp = torch.exp(x-means)
    x_exp_sum = torch.sum(x_exp, 1, keepdim=True)

    return x_exp/x_exp_sum

and found that this implementation can achieve better accuracy. However, it is still not as good as the F.softmax. Anyone can help?