[BUG] Training with activation function 'gelu' on CPU results in nan

[liangadam]

Summary

In a SE_A descriptor's example, when we choose the activation function of the descriptor in input.json to be 'gelu' and train it on CPU, we will get 'nan' in lcurve and in dp test's result.

Deepmd-kit version, installation way, input file, running commands, error log, etc.

version: v2.0.0-beta.4 (devel branch)
installation way: build python interface(tensorflow) and pip install deepmd-kit in a conda environment.
input file: example/water/se_e2_a/input.json ,and specify the activation function of descriptor is gelu.
running command: dp train input.json
error log: nan in lcurve
detail: run it on CPU, without command 'export DP_VARIANT=cuda'

platform: Lebesgue base image

Steps to Reproduce
In example/water/se_e2_a open input.json ,and specify the activation function of descriptor is gelu, then run 'dp train input.json' on CPU.

Further Information, Files, and Links

[njzjz]

I can reproduce nan, only when GPU is not avaiable.

[njzjz]

And this only happened when training forces.

[njzjz]

DEEPMD INFO batch 2 training time 0.25 s, testing time 0.13 s
gelu_grad_cpu out[147200] is -nan, dy[147200] is -nan, xx[147200] is -0.345248.
gelu_grad_cpu out[73600] is -nan, dy[73600] is -nan, xx[73600] is -1.090311.
gelu_cpu out[0] is -nan, xx[0] is -nan.
gelu_cpu out[0] is -nan, xx[0] is -nan.
gelu_cpu out[0] is -nan, xx[0] is -nan.
gelu_grad_cpu out[0] is -nan, dy[0] is -nan, xx[0] is 0.076259.
gelu_grad_cpu out[0] is -nan, dy[0] is -nan, xx[0] is -nan.
gelu_grad_cpu out[0] is -nan, dy[0] is -nan, xx[0] is -nan.
gelu_grad_cpu out[0] is -nan, dy[0] is -nan, xx[0] is -nan.
gelu_grad_cpu out[0] is -nan, dy[0] is -nan, xx[0] is 0.562668.
gelu_grad_cpu out[0] is -nan, dy[0] is -nan, xx[0] is 0.779181.

[Ericwang6]

I met the same problem when training on my system with gelu function, and this training is on GPU. The version of deepmd-kit is 2.0.0.b3 installed from off-line package.
Attached is the data to reproduce the issue.
issue.tar.gz

[denghuilu]

#939 should fix this problem.