GPU memory leak when using soft_sphere_neighbor_list with epsilon species tensor

[ESEberhard]

I have two species in my simulation.
I wanted to use a 2x2 epsilon matrix in the same way that I am using a 2x2 sigma matrix

sigma = jnp.array([ [1e-7,                       params['sigma_mem_enz'] ],
                    [params['sigma_mem_enz'],    params['sigma-enz-enz'] ]], dtype=jnp.float32)

epsilon = jnp.array([ [0                        ,  params['epsilon_mem_enz'] ],
                       [params['epsilon_mem_enz'],  params['epsilon-enz-enz'] ]], dtype=jnp.float32)

neighbor_fn, enz_mem_pot = jax_md.energy.soft_sphere_neighbor_list(
    displacement_fn,
    BOX_SIZE,
    species,
    sigma,
    epsilon,
    alpha=2,
    dr_threshold=NEIGHBOR_LIST_DR_THRESHOLD,
)

But this slowly leaks memory.
I can observe the GPU memory slowly going up until it crashes after a couple of hours.
If I instead use a regular float for epsilon I do not observe this behaviour.
The memory persists when using onp arrays instead too