[content] /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/causal/dynamics-cotraining/train_dynamics_causal_8_node_3412356.log (log)
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
2025-08-10 15:25:51.077766: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_1_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.083044: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.083756 2894266 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
2025-08-10 15:25:51.086970: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_1_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.092551: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.093334  980422 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:51.201756: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_1_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.206590: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.207293  634191 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 14, which is shorter than the requested sequence length 16.
2025-08-10 15:25:51.396407: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.401620: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.402306 1432779 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:51.439957: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_2_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.444801: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.445487 1432781 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:51.469245: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.474418: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
2025-08-10 15:25:51.474187: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_1_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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.
E0810 15:25:51.475149  980421 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
2025-08-10 15:25:51.479062: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.479774 2116889 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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.
2025-08-10 15:25:51.480531: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:51.486278: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.486998 2894265 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
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:

2025-08-10 15:25:51.494490: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_2_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.499746: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.500327 1675361 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
Traceback (most recent call last):
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:51.539849: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_1_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.540027: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.545326: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
2025-08-10 15:25:51.545615: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.546093  370067 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
E0810 15:25:51.546508  486006 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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:

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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Traceback (most recent call last):
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 10, which is shorter than the requested sequence length 16.
Filtering out episode with length 8, which is shorter than the requested sequence length 16.
2025-08-10 15:25:51.644759: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_1_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.649863: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.650467 1432780 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:51.755230: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_2_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.760070: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.760746  980423 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 4, which is shorter than the requested sequence length 16.
Filtering out episode with length 15, which is shorter than the requested sequence length 16.
Filtering out episode with length 9, which is shorter than the requested sequence length 16.
Filtering out episode with length 15, which is shorter than the requested sequence length 16.
Filtering out episode with length 4, which is shorter than the requested sequence length 16.
Filtering out episode with length 12, which is shorter than the requested sequence length 16.
Filtering out episode with length 14, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
2025-08-10 15:25:51.928669: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_2_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.933969: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.934606  634192 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 14, which is shorter than the requested sequence length 16.
2025-08-10 15:25:51.969201: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_2_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:51.974909: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:51.975612 2894267 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 12, which is shorter than the requested sequence length 16.
2025-08-10 15:25:52.058653: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_2_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.064168: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.064814  486007 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
2025-08-10 15:25:52.069997: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_1_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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.
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:52.075172: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.075840  370068 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 12, which is shorter than the requested sequence length 16.
Filtering out episode with length 6, which is shorter than the requested sequence length 16.
2025-08-10 15:25:52.164898: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.169712: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.170418  634190 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
Filtering out episode with length 12, which is shorter than the requested sequence length 16.
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:52.177594: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_3_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.182418: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.183116  370070 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 15, which is shorter than the requested sequence length 16.
Filtering out episode with length 14, which is shorter than the requested sequence length 16.
Filtering out episode with length 2, which is shorter than the requested sequence length 16.
Filtering out episode with length 12, which is shorter than the requested sequence length 16.
2025-08-10 15:25:52.533630: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_3_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.538658: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.539439 1432782 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
2025-08-10 15:25:52.544122: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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.
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
2025-08-10 15:25:52.544952: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_2_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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.
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
2025-08-10 15:25:52.548998: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
2025-08-10 15:25:52.549898: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.549718 1675359 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
E0810 15:25:52.550622  370069 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:52.561947: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_3_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.566937: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.567615 2894268 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:52.642384: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_3_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.647302: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.647957 2116891 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 4, which is shorter than the requested sequence length 16.
2025-08-10 15:25:52.745959: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_3_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.750851: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.751511  980424 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:52.759404: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_3_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.764414: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.765291  634193 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:52.780996: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_2_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.786535: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.787190 2116890 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 2, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
2025-08-10 15:25:52.978366: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:52.983245: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:52.983873 2116888 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:53.011901: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_3_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:53.016696: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:53.017413  486008 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
2025-08-10 15:25:53.058233: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:53.063123: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:53.063822  486005 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 15, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 3, which is shorter than the requested sequence length 16.
2025-08-10 15:25:53.209904: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_3_bfc) ran out of memory trying to allocate 16.04GiB (rounded to 17224249088)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-08-10 15:25:53.215211: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] **************************************************__________________________________________________
E0810 15:25:53.215882 1675362 pjrt_stream_executor_client.cc:2916] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes. [tf-allocator-allocation-error='']
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 "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/train_dynamics.py", line 365, in <module>
    loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
  File "/hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar/.venv/lib/python3.10/site-packages/flax/nnx/transforms/compilation.py", line 431, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
jaxlib._jax.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 17224248880 bytes.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 6, which is shorter than the requested sequence length 16.
Filtering out episode with length 6, which is shorter than the requested sequence length 16.
Filtering out episode with length 12, which is shorter than the requested sequence length 16.
Filtering out episode with length 6, which is shorter than the requested sequence length 16.
Filtering out episode with length 11, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 12, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 13, which is shorter than the requested sequence length 16.
Filtering out episode with length 2, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
Filtering out episode with length 13, which is shorter than the requested sequence length 16.
Filtering out episode with length 1, which is shorter than the requested sequence length 16.
wandb: 
wandb: 🚀 View run dynamics-causal-8-node-3412356 at: https://wandb.ai/instant-uv/jafar/runs/3412356
wandb: Find logs at: ../../../../../hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jafar_jobs/wandb/run-20250810_152243-3412356/logs
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 11 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/home/hk-project-p0023960/tum_cte0515/.local/share/uv/python/cpython-3.10.18-linux-x86_64-gnu/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 10 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
srun: error: hkn0527: task 29: Exited with exit code 1
srun: error: hkn0522: task 10: Exited with exit code 1
srun: error: hkn0523: task 13: Exited with exit code 1
srun: error: hkn0404: tasks 1-2: Exited with exit code 1
srun: error: hkn0524: tasks 17-18: Exited with exit code 1
srun: error: hkn0525: tasks 20,22: Exited with exit code 1
srun: error: hkn0527: tasks 28,30: Exited with exit code 1
srun: error: hkn0522: tasks 8-9: Exited with exit code 1
srun: error: hkn0523: tasks 14-15: Exited with exit code 1
srun: error: hkn0521: tasks 4-6: Exited with exit code 1
srun: error: hkn0526: tasks 24-26: Exited with exit code 1
srun: error: hkn0527: task 31: Exited with exit code 1
srun: error: hkn0522: task 11: Exited with exit code 1
srun: error: hkn0521: task 7: Exited with exit code 1
srun: error: hkn0404: task 3: Exited with exit code 1
srun: error: hkn0524: task 19: Exited with exit code 1
srun: error: hkn0525: tasks 21,23: Exited with exit code 1
srun: error: hkn0523: task 12: Exited with exit code 1
srun: error: hkn0524: task 16: Exited with exit code 1
srun: error: hkn0404: task 0: Exited with exit code 1
srun: error: hkn0526: task 27: Exited with exit code 1

============================= JOB FEEDBACK =============================

Job ID: 3412356
Cluster: hk
User/Group: tum_cte0515/hk-project-p0023960
Account: hk-project-p0023960
State: FAILED (exit code 1)
Partition: accelerated
Nodes: 8
Cores per node: 24
Nodelist: hkn[0404,0521-0527]
CPU Utilized: 01:45:59
CPU Efficiency: 13.30% of 13:16:48 core-walltime
Job Wall-clock time: 00:04:09
Starttime: Sun Aug 10 15:21:49 2025
Endtime: Sun Aug 10 15:25:58 2025
Memory Utilized: 234.69 GB (estimated maximum)
Memory Efficiency: 0.00% of 0.00 MB (0.00 MB/node)
Energy Consumed: 13898475 Joule / 3860.6875 Watthours
Average node power draw: 55817.1686746988 Watt