[insane_dreamer]

Yes because the pypi cupy/cudann packages now work seamlessly with jax. Until not long ago we had to use the conda packages.