Jax GPU Resouce Exhausted [Preallocate =False did not work]

[tan-liam]

I have looked around, but couldn't find any working solutions. Most of the posts recommend setting XLA_PYTHON_CLIENT_PREALLOCATE=false
But this did not work for me.

I am running a very small model.

50 hidden layers, 100 units each layer, with the input being size 25000. Batch size 32.

This should not be over 100 mb.

Does anyone know of any newer fix?

RESOURCE_EXHAUSTED: Out of memory while trying to allocate 25000000256 bytes

Read another post saying this is a user Error:

However I have a hard time believing that the program requires 25gb to run.
It is a very small fully connected network.