CUDA error when using meshio

Im testing JAX-FEM in a newer machine running WSL on Windows 11, with a RTX A4500 20Gb GPU. Everything appears to be up to date, including the GPU drivers.

The Windows host is using CUDA 13.3.

When I try to run a very small model (fewer than 5,000 nodes), I get the following error.

`E0630 11:22:47.221764 442 cuda_executor.cc:1182] [0] Failed to allocate device memory of 20.8KiB (21248 bytes): INTERNAL: CUDA error: : CUDA_ERROR_UNKNOWN: unknown error
E0630 11:22:47.222657 442 cuda_executor.cc:1182] [0] Failed to allocate device memory of 18.8KiB (19200 bytes): INTERNAL: CUDA error: : CUDA_ERROR_UNKNOWN: unknown error
E0630 11:22:47.223638 442 cuda_executor.cc:1182] [0] Failed to allocate device memory of 17.0KiB (17408 bytes): INTERNAL: CUDA error: : CUDA_ERROR_UNKNOWN: unknown error
W0630 11:22:47.223656 442 bfc_allocator.cc:514] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.2KiB (rounded to 16640)requested by op
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation.
Current allocation summary follows.
Current allocation summary follows.
W0630 11:22:47.223726 442 bfc_allocator.cc:525]
Traceback (most recent call last):
File "/home/xxxxxx/main.py", line 384, in
main()
File "/home/xxxxxx/main.py", line 371, in main
mesh, elem_layers, ele_type = build_box_mesh(cfg)
^^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/main.py", line 83, in build_box_mesh
mesh = Mesh(meshio_mesh.points,meshio_mesh.cells_dict[cell_type])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax_fem/generate_mesh.py", line 25, in __init__
self.points = np.asarray(points)
^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/numpy/array_constructors.py", line 453, in asarray
return array(a, dtype=dtype, copy=bool(copy), order=order, device=device,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/numpy/array_constructors.py", line 208, in array
return lax.stage(object)
^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/lax/lax.py", line 1626, in stage
return stage_p.bind(x)
^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/core.py", line 705, in bind
return self.bind_with_trace(prev_trace, args, avals, params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/lax/lax.py", line 5406, in _stage_bind_with_trace
return trace.stage_value(args[0])
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/core.py", line 1283, in stage_value
return self.process_primitive(stage_p, [val], {})
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/core.py", line 1294, in process_primitive
return primitive.impl(*args, **params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/lax/lax.py", line 5396, in _stage_impl
return dispatch.apply_primitive(stage_p, x)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/xxxxxx/miniconda3/envs/testepy3.12.1/lib/python3.12/site-packages/jax/_src/dispatch.py", line 88, in apply_primitive
outs = fun(*args)
^^^^^^^^^^
jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.2KiB. [tf-allocator-allocation-error='']
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.`

Initially, I was using Python 3.13, where I encountered this error. To rule out a compatibility issue, I created a new environment with Python 3.12 and CUDA 13.1, but the same error persists.

The only related information I found is that JAX sometimes preallocates GPU memory, which can lead to similar errors. I disabled GPU memory preallocation, but the issue remains unchanged.

Has anyone encountered this type of error when using JAX-FEM? Any suggestions on what else I should check?

For comparison, I have an older JAX-FEM installation running on another PC (also using WSL, but on Windows 10 with CUDA 12), and it runs smoothly. Because of that, WSL itself does not appear to be the root cause.