fix: removed the jax import and defaulted jax to use the cpu due to out of memory issues in the unet demo - Githubissues

ivy-llc / demos

Demos and Tutorials with Ivy

https://ivy.dev/docs/demos/

19 stars 61 forks source link

fix: removed the jax import and defaulted jax to use the cpu due to out of memory issues in the unet demo #72

Closed vedpatwardhan closed 10 months ago

vedpatwardhan commented 10 months ago

The jax import issue is mentioned here : https://github.com/google/jax/issues/18032