Errors raised in the function: utils.render_image (nerf_sh.train, line 255)

[TanGeeGo]

Thanks for the contribution of the authors.
In the first 20000 iterations, the provided codes run normally. But in the test set evaluation, errors of utils.render_image stops the evaluation. Specific errors like this:

File "/hdd8T_4/chensq/PlenOctree/nerf_sh/train.py", line 255, in main
pred_color, pred_disp, pred_acc = utils.render_image(
File "/hdd8T_4/chensq/PlenOctree/nerf_sh/nerf/utils.py", line 350, in render_image
unusedrng, key_0, key_1 = jax.random.split(rng, 3)
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/_src/random.py", line 270, in split
return _split(key, int(num)) # type: ignore
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 143, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/_src/api.py", line 426, in cache_miss
out_flat = xla.xla_call(
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/core.py", line 1565, in bind
return call_bind(self, fun, *args, **params)
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/core.py", line 1556, in call_bind
outs = primitive.process(top_trace, fun, tracers, params)
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/core.py", line 1568, in process
return trace.process_call(self, fun, tracers, params)
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/core.py", line 609, in process_call
return primitive.impl(f, *tracers, **params)
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/interpreters/xla.py", line 581, in _xla_call_impl
return compiled_fun(*args)
File "/hdd8T_4/chensq/anaconda3/envs/plenoctree/lib/python3.8/site-packages/jax/interpreters/xla.py", line 874, in _execute_compiled
out_bufs = compiled.execute(input_bufs)
jax._src.traceback_util.UnfilteredStackTrace: RuntimeError: CUDA operation failed: invalid resource handle

There is something wrong with the function jax.random.split(rng, 3). I don't know how to resolve it.
Looking forward to the authors' reply!