Flax Quick Start Tutorial requires 7GB of VRAM ?! help please! #4082
Replies: 1 comment
-
|
Ok i guess this explains it: https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html. Maybe would be good to mention it in the Quick Start. |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Hi,
i tried out the Flax Quick Start Tutorial without any modification and the VRAM usage is 7GB which can not be correct. I benchmarked a similar example and the same issue occured. In the similar exampe I tested against pytorch which used 300mb vs 6GB for an Autoencoder on Mnist. Not sure what's going on. I couldn't find anything on the internet.
Beta Was this translation helpful? Give feedback.
All reactions