Thanks for replying !!
Can autograd.gradcheck() detect memory leakage ? One of my problem is that the result is normal when I use only my cude implementation. But when I use both my cuda implementation and pytorch autograd to check the difference, in the cuda implementation branch, the backward gradient of layers that uses pytorch inner autograd  becomes very large and finally gets nan after 5 iters, even though the gradient of my cuda implementation has identical backward gradient if the input logits and labels are same. Thus I suspect that there might be some memory leakage that changed the values of other layer parameters.