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.