AttributeError: module 'jax.random' has no attribute 'KeyArray'

real-stanford / diffusion_policy

[RSS 2023] Diffusion Policy Visuomotor Policy Learning via Action Diffusion

https://diffusion-policy.cs.columbia.edu/

MIT License

1.09k stars 203 forks source link

AttributeError: module 'jax.random' has no attribute 'KeyArray' #67

Open BE-A-HERO opened 2 months ago

BE-A-HERO commented 2 months ago

Hi! Recently, when I tried to conduct the example code diffusion_policy_state_pusht_demo.ipynb and diffusion_policy_vision_pusht_demo.ipynb, it returns error: 'AttributeError: module 'jax.random' has no attribute 'KeyArray'.

It seems like the version problem. Are there any solutions? Thanks.

JohnZhan2023 commented 2 months ago

!pip3 install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html may help