[terminal_command] TERMINAL ()
    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1256: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.  warnings.warn(2025-07-25 16:22:58.570793: W external/xla/xla/service/platform_util.cc:220] unable to create StreamExecutor for CUDA:0: : CUDA_ERROR_OUT_OF_MEMORY: out of memory2025-07-25 16:22:58.583791: W external/xla/xla/service/platform_util.cc:220] unable to create StreamExecutor for CUDA:0: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryTraceback (most recent call last):  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 742, in backends    backend = _init_backend(platform)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 828, in _init_backend    backend = registration.factory()  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 518, in factory    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)Traceback (most recent call last):  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 742, in backends    backend = _init_backend(platform)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 828, in _init_backend    backend = registration.factory()  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 518, in factory    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)2025-07-25 16:22:58.786439: W external/xla/xla/service/platform_util.cc:220] unable to create StreamExecutor for CUDA:0: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryTraceback (most recent call last):  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 742, in backends    backend = _init_backend(platform)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 828, in _init_backend    backend = registration.factory()  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 518, in factory    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)2025-07-25 16:23:00.095031: W external/xla/xla/service/platform_util.cc:220] unable to create StreamExecutor for CUDA:0: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryTraceback (most recent call last):  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 742, in backends    backend = _init_backend(platform)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 828, in _init_backend    backend = registration.factory()  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 518, in factory    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 174, in <module>    video_batch = next(iter(dataloader))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/grain/_src/python/data_loader.py", line 479, in __next__    result_record = next(self._iterator)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/grain/_src/python/data_loader.py", line 380, in _iterator_with_context    yield from it  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/grain/_src/python/grain_pool.py", line 784, in __next__    raise element  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/grain/_src/python/grain_pool.py", line 607, in _process_elements_in_grain_pool    with GrainPool(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/grain/_src/python/grain_pool.py", line 422, in __init__    parallel.run_in_parallel(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/grain/_src/core/parallel.py", line 70, in run_in_parallel    raise completed.exception()  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/concurrent/futures/thread.py", line 58, in run    result = self.fn(*self.args, **self.kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/grain/_src/python/grain_pool.py", line 423, in <lambda>    function=lambda child_process: child_process.start(),  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/process.py", line 121, in start    self._popen = self._Popen(self)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/context.py", line 288, in _Popen    return Popen(process_obj)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__    super().__init__(process_obj)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__    self._launch(process_obj)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 42, in _launch    prep_data = spawn.get_preparation_data(process_obj._name)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 154, in get_preparation_data    _check_not_importing_main()  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 134, in _check_not_importing_main    raise RuntimeError('''RuntimeError:         An attempt has been made to start a new process before the        current process has finished its bootstrapping phase.        This probably means that you are not using fork to start your        child processes and you have forgotten to use the proper idiom        in the main module:            if __name__ == '__main__':                freeze_support()                ...        The "freeze_support()" line can be omitted if the program        is not going to be frozen to produce an executable.(jafar) [franz.srambical@hai003.haicore.berlin:~/jafar] $ bash experiments/sample.sh /fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1256: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.  warnings.warn(E0725 16:24:15.264406 1600031 cuda_dnn.cc:535] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERRORE0725 16:24:15.264473 1600031 cuda_dnn.cc:539] Memory usage: 131072 bytes free, 85029158912 bytes total.E0725 16:24:15.267300 1600031 cuda_dnn.cc:535] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERRORE0725 16:24:15.267338 1600031 cuda_dnn.cc:539] Memory usage: 131072 bytes free, 85029158912 bytes total.E0725 16:24:15.269648 1600037 cuda_dnn.cc:535] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERRORE0725 16:24:15.269799 1600037 cuda_dnn.cc:539] Memory usage: 131072 bytes free, 85029158912 bytes total.E0725 16:24:15.272470 1600037 cuda_dnn.cc:535] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERRORE0725 16:24:15.278304 1600037 cuda_dnn.cc:539] Memory usage: 131072 bytes free, 85029158912 bytes total.jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Bad StatusOr access: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryE0725 16:24:15.397122 1600031 cuda_dnn.cc:535] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERRORE0725 16:24:15.404450 1600031 cuda_dnn.cc:539] Memory usage: 131072 bytes free, 85029158912 bytes total.E0725 16:24:15.407978 1600031 cuda_dnn.cc:535] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERRORE0725 16:24:15.408150 1600031 cuda_dnn.cc:539] Memory usage: 131072 bytes free, 85029158912 bytes total.jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)jaxlib._jax.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.E0725 16:24:15.478617 1600037 cuda_dnn.cc:535] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERRORE0725 16:24:15.493273 1600037 cuda_dnn.cc:539] Memory usage: 131072 bytes free, 85029158912 bytes total.E0725 16:24:15.497270 1600037 cuda_dnn.cc:535] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERRORE0725 16:24:15.497309 1600037 cuda_dnn.cc:539] Memory usage: 131072 bytes free, 85029158912 bytes total.jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)jaxlib._jax.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Bad StatusOr access: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryjax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Bad StatusOr access: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryERROR:absl:Processing Failed. Shutting down./fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1256: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.  warnings.warn(/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1256: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.  warnings.warn(2025-07-25 16:24:37.987107: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 4.88MiB (rounded to 5120000)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.2025-07-25 16:24:37.987805: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] *******x********************************************************************************************E0725 16:24:37.987942 1600038 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 5120000 bytes. [tf-allocator-allocation-error='']Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 68, in <module>    genie = Genie(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 143, in __call__    return _graph_node_meta_call(cls, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 152, in _graph_node_meta_call    cls._object_meta_construct(node, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 146, in _object_meta_construct    self.__init__(*args, **kwargs)  File "/fast/home/franz.srambical/jafar/genie.py", line 94, in __init__    self.lam = LatentActionModel(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 143, in __call__    return _graph_node_meta_call(cls, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 152, in _graph_node_meta_call    cls._object_meta_construct(node, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 146, in _object_meta_construct    self.__init__(*args, **kwargs)  File "/fast/home/franz.srambical/jafar/models/lam.py", line 83, in __init__    self.decoder = STTransformer(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 143, in __call__    return _graph_node_meta_call(cls, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 152, in _graph_node_meta_call    cls._object_meta_construct(node, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 146, in _object_meta_construct    self.__init__(*args, **kwargs)  File "/fast/home/franz.srambical/jafar/utils/nn.py", line 191, in __init__    STBlock(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 143, in __call__    return _graph_node_meta_call(cls, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 152, in _graph_node_meta_call    cls._object_meta_construct(node, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 146, in _object_meta_construct    self.__init__(*args, **kwargs)  File "/fast/home/franz.srambical/jafar/utils/nn.py", line 51, in __init__    self.spatial_pos_enc = PositionalEncoding(self.dim)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 143, in __call__    return _graph_node_meta_call(cls, *args, **kwargs)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/flax/nnx/object.py", line 152, in _graph_node_meta_call(jafar) [franz.srambical@hai003.haicore.berlin:~/jafar] $ bash experiments/sample.sh ^Csrun: interrupt (one more within 1 sec to abort)srun: StepId=13885.1 task 0: running^Csrun: sending Ctrl-C to StepId=13885.1srun: forcing job terminationsrun: Job step aborted: Waiting up to 32 seconds for job step to finish.[2025-07-25T16:32:05.412] error: *** STEP 13885.1 ON hai003 CANCELLED AT 2025-07-25T16:32:05 DUE to SIGNAL Killed ***^Csrun: sending Ctrl-C to StepId=13885.1srun: job abort in progress(jafar) [franz.srambical@hai003.haicore.berlin:~/jafar] $ bash experiments/sample.sh /fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1256: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.  warnings.warn(jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Bad StatusOr access: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory2025-07-25 16:32:55.422761: W external/xla/xla/service/platform_util.cc:220] unable to create StreamExecutor for CUDA:0: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryTraceback (most recent call last):  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 742, in backends    backend = _init_backend(platform)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 828, in _init_backend    backend = registration.factory()  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 518, in factory    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)2025-07-25 16:32:55.572021: W external/xla/xla/service/platform_util.cc:220] unable to create StreamExecutor for CUDA:0: : CUDA_ERROR_OUT_OF_MEMORY: out of memory2025-07-25 16:32:55.575094: W external/xla/xla/service/platform_util.cc:220] unable to create StreamExecutor for CUDA:0: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryTraceback (most recent call last):  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 742, in backends    backend = _init_backend(platform)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 828, in _init_backend    backend = registration.factory()  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 518, in factory    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)Traceback (most recent call last):  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 742, in backends    backend = _init_backend(platform)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 828, in _init_backend    backend = registration.factory()  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 518, in factory    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)2025-07-25 16:32:55.985487: W external/xla/xla/service/platform_util.cc:220] unable to create StreamExecutor for CUDA:0: : CUDA_ERROR_OUT_OF_MEMORY: out of memoryTraceback (most recent call last):  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 742, in backends    backend = _init_backend(platform)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 828, in _init_backend    backend = registration.factory()  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/xla_bridge.py", line 518, in factory    return xla_client.make_c_api_client(plugin_name, updated_options, None)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jaxlib/xla_client.py", line 153, in make_c_api_client    return _xla.get_c_api_client(plugin_name, options, distributed_client)jaxlib._jax.XlaRuntimeError: INTERNAL: no supported devices found for platform CUDADuring handling of the above exception, another exception occurred:jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.The above exception was the direct cause of the following exception:Traceback (most recent call last):  File "<string>", line 1, in <module>  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main    exitcode = _main(fd, parent_sentinel)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 125, in _main    prepare(preparation_data)  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 236, in prepare    _fixup_main_from_path(data['init_main_from_path'])  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/spawn.py", line 287, in _fixup_main_from_path    main_content = runpy.run_path(main_path,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 289, in run_path    return _run_module_code(code, init_globals, run_name,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 96, in _run_module_code    _run_code(code, mod_globals, init_globals,  File "/home/franz.srambical/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/runpy.py", line 86, in _run_code    exec(code, run_globals)  File "/fast/home/franz.srambical/jafar/sample.py", line 64, in <module>    rng = jax.random.PRNGKey(args.seed)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 249, in PRNGKey    return _return_prng_keys(True, _key('PRNGKey', seed, impl))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/random.py", line 201, in _key    return prng.random_seed(seed, impl=impl)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 551, in random_seed    seeds_arr = jnp.asarray(np.int64(seeds))  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 383, in asarray    return array(a, dtype=dtype, copy=bool(copy), order=order, device=device)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/numpy/array.py", line 268, in array    out_array: Array = lax._convert_element_type(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 1726, in _convert_element_type    return convert_element_type_p.bind(  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 536, in bind    return self._true_bind(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 552, in _true_bind    return self.bind_with_trace(prev_trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 4902, in _convert_element_type_bind_with_trace    operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 562, in bind_with_trace    return trace.process_primitive(self, args, params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 1066, in process_primitive    return primitive.impl(*args, **params)  File "/fast/home/franz.srambical/jafar/.venv/lib/python3.10/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive    outs = fun(*args)RuntimeError: Unable to initialize backend 'cuda': INTERNAL: no supported devices found for platform CUDA (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)^Csrun: interrupt (one more within 1 sec to abort)srun: StepId=13885.2 task 0: running^Csrun: sending Ctrl-C to StepId=13885.2srun: forcing job terminationsrun: Job step aborted: Waiting up to 32 seconds for job step to finish.[2025-07-25T16:33:04.884] error: *** STEP 13885.2 ON hai003 CANCELLED AT 2025-07-25T16:33:04 DUE to SIGNAL Killed ***^Csrun: sending Ctrl-C to StepId=13885.2srun: job abort in progress(jafar) [franz.srambical@hai003.haicore.berlin:~/jafar] $ ^C(jafar) [franz.srambical@hai003.haicore.berlin:~/jafar] $ nvidia-smi Fri Jul 25 16:33:07 2025       +-----------------------------------------------------------------------------------------+| NVIDIA-SMI 570.124.06             Driver Version: 570.124.06     CUDA Version: 12.8     ||-----------------------------------------+------------------------+----------------------+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC || Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. ||                                         |                        |               MIG M. ||=========================================+========================+======================||   0  NVIDIA H100 80GB HBM3          On  |   00000000:19:00.0 Off |                    0 || N/A   42C    P0            128W /  700W |       1MiB /  81559MiB |      0%      Default ||                                         |                        |             Disabled |+-----------------------------------------+------------------------+----------------------+                                                                                         +-----------------------------------------------------------------------------------------+| Processes:                                                                              ||  GPU   GI   CI              PID   Type   Process name                        GPU Memory ||        ID   ID                                                               Usage      ||=========================================================================================||  No running processes found                                                             |+-----------------------------------------------------------------------------------------+(jafar) [franz.srambical@hai003.haicore.berlin:~/jafar] $ exitexitsalloc: Relinquishing job allocation 13885