search

real-stanford / diffusion_policy

[RSS 2023] Diffusion Policy Visuomotor Policy Learning via Action Diffusion
https://diffusion-policy.cs.columbia.edu/
MIT License
1.61k stars 321 forks source link

module 'jax.random' has no attribute 'KeyArray' #105

Open monk279 opened 2 months ago

monk279 commented 2 months ago

"module 'jax.random' has no attribute 'KeyArray'" occurs at the second block in both two Google Colab notebooks. It seems like the something goes wrong with the diffuerser from huggingface.

denisgriaznov commented 2 months ago

I faced the same problem, found a solution, came here to ask the developers for changes and saw your question :) This problem is solved by downgrading jax:

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