Train script doesnt run about plenoctree

[grgkopanas]

I also tried to manully install all dependencies by following the guidelines of each respective pacakage (jax, flax, tensorflow etc)
and ended up in another error which is deeper in the code:

I0119 11:16:00.161019 140446576289600 checkpoints.py:249] Found no checkpoint files in ckpts/chair
/home/gkopanas/.conda/envs/plenoctrees4/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py:413: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.
  warnings.warn(
* Prefetch
/home/gkopanas/.conda/envs/plenoctrees4/lib/python3.8/site-packages/jax/interpreters/xla.py:799: UserWarning: Some donated buffers were not usable: f32[1024,3]{1,0}, f32[1024,3]{1,0}, f32[1024,3]{1,0}, f32[1024,3]{1,0}
  warnings.warn("Some donated buffers were not usable: {}".format(
Traceback (most recent call last):
  File "/home/gkopanas/.conda/envs/plenoctrees4/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/gkopanas/.conda/envs/plenoctrees4/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/data/graphdeco/user/gkopanas/plenoctree/nerf_sh/train.py", line 314, in <module>
    app.run(main)
  File "/home/gkopanas/.conda/envs/plenoctrees4/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/gkopanas/.conda/envs/plenoctrees4/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "/data/graphdeco/user/gkopanas/plenoctree/nerf_sh/train.py", line 198, in main
    state, stats, keys = train_pstep(keys, state, batch, lr)
ValueError: INTERNAL: Failed to launch CUDA kernel: fusion_201 with block dimensions: 1024x1x1 and grid dimensions: 576x1x1: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered

[Holmes-Alan]

Have you resolved the problem?

[drcfts]

Were you or @Holmes-Alan able to solve this problem?

[chenyuntc]

The key is to setup jax env