Error: Fatal Python error: Segmentation fault

[mileywang3061]

I0414 10:36:42.831109 140578513020736 train.py:273] Starting training
Current thread 0x00007fed4e792740 (most recent call first):
File "/home/anaconda3/envs/nerfies/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 2033 in shape
File "<array_function internals>", line 200 in shape
File "/home/anaconda3/envs/nerfies/lib/python3.8/site-packages/jax/_src/api.py", line 1307 in
File "/home/anaconda3/envs/nerfies/lib/python3.8/site-packages/jax/_src/api.py", line 1307 in _mapped_axis_size
File "/home/anaconda3/envs/nerfies/lib/python3.8/site-packages/jax/_src/api.py", line 1633 in f_pmapped
File "/home/anaconda3/envs/nerfies/lib/python3.8/site-packages/jax/_src/api.py", line 1725 in f_pmapped
File "/home/anaconda3/envs/nerfies/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162 in reraise_with_filtered_traceback
File "train.py", line 296 in main
File "/home/anaconda3/envs/nerfies/lib/python3.8/site-packages/absl/app.py", line 254 in _run_main
File "/home/anaconda3/envs/nerfies/lib/python3.8/site-packages/absl/app.py", line 308 in run
File "train.py", line 334 in
srun: error: storm: task 0: Segmentation fault

Details: tensorflow=2.7.0
jax=0.2.20
jaxlib=0.3.2
numpy=1.19.2

Please give me some advice how to handle this problem.