JIT session error: Cannot allocate memory

[JuanFMontesinos]

Hi,
I was using jax and training a neural network for time series event detection.
In essence, at some point I have a non-jited function that calls the following code:

try:
    slice_data = softmax_prob[cc.start : cc.end, i]
except:
    print(f"Error: {cc.start}, {cc.end}, {i}, {softmax_prob.shape}")
    breakpoint()

where the try-except was added for debugging purposed. When trying to slice that array, it raises the JIT session error: Cannot allocate memory.
With the debugger, I could check that, for that particular example, the values seems overall right:

If I try to slice (once in the debugger), I obtain this

I was also carefully following the system stats, there was enough space in tmp, more than 100 Gb RAM free, and 6 Gb GPU RAM still free too. Also this happened in after training iteration 8000+ and after testing iteration 800, which means it's an exotic error.

Do anybody have any clue what could be the error?

[Qazalbash]

I had a similar issue. It was resolved by setting XLA_PYTHON_CLIENT_ALLOCATOR=platform and changing some hyperparameters.