/home/dat/pino/lib/python3.8/site-packages/jax/lib/xla_bridge.py:386: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.
  warnings.warn(
/home/dat/pino/lib/python3.8/site-packages/jax/lib/xla_bridge.py:373: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.
  warnings.warn(
Epoch ... (1/3):   0%|                                                                                                        | 0/3 [00:00<?, ?it/s]2021-07-14 19:10:58.993859: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2036] Execution of replica 0 failed: Resource exhausted: Attempting to allocate 17.0K. That was not possible. There are 48.0K free. Due to fragmentation, the largest contiguous region of free memory is 16.0K.; (0x0x0_HBM0)
Epoch ... (1/3):   0%|                                                                                                        | 0/3 [15:02<?, ?it/s]
Traceback (most recent call last):
  File "./run_mlm_flax_no_accum.py", line 686, in <module>
    train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)
  File "./run_mlm_flax_no_accum.py", line 255, in generate_batch_splits
    batch_idx = np.split(samples_idx, sections_split)
  File "<__array_function__ internals>", line 5, in split
  File "/home/dat/pino/lib/python3.8/site-packages/numpy/lib/shape_base.py", line 874, in split
    return array_split(ary, indices_or_sections, axis)
  File "<__array_function__ internals>", line 5, in array_split
  File "/home/dat/pino/lib/python3.8/site-packages/numpy/lib/shape_base.py", line 790, in array_split
    sub_arys.append(_nx.swapaxes(sary[st:end], axis, 0))
  File "/home/dat/pino/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 5009, in _rewriting_take
    return _gather(arr, treedef, static_idx, dynamic_idx, indices_are_sorted,
  File "/home/dat/pino/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 5028, in _gather
    y = lax.gather(
  File "/home/dat/pino/lib/python3.8/site-packages/jax/_src/lax/lax.py", line 984, in gather
    return gather_p.bind(
  File "/home/dat/pino/lib/python3.8/site-packages/jax/core.py", line 264, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/core.py", line 603, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/interpreters/xla.py", line 249, in apply_primitive
    return compiled_fun(*args)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/interpreters/xla.py", line 365, in _execute_compiled_primitive
    out_bufs = compiled.execute(input_bufs)
RuntimeError: Resource exhausted: Attempting to allocate 17.0K. That was not possible. There are 48.0K free. Due to fragmentation, the largest contiguous region of free memory is 16.0K.; (0x0x0_HBM0)