Hi Ming,
Try to update your NVIDIA drivers to latest version.
This is definitely an issue either related to the nvidia driver or related to hardware issue (overheating or other issue).

We also tried to use MXNET train on 4 GPUs, it goes well without seeing this problem.

This could be because when you are using MXNet, the mxnet install did not end up using cudnn or nccl (for example) and runs fine.
I would also suggest that you try installing pytorch from source or docker image, just to see if that helps your case. It can also help me debug your case: https://github.com/pytorch/pytorch#from-source 4