RESOURCE_EXHAUSTED: Out of memory while trying to allocate * bytes.

[devavratTomar]

Hi,

I was trying jaxline with 2 local GPUs, but I receive RESOURCE_EXHAUSTED: Out of memory while trying to allocate 285212672 bytes. errors. I have sufficient memory available on the GPUs.

Thank you.