/home/dat/pino/lib/python3.8/site-packages/jax/lib/xla_bridge.py:382: 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:369: 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/5):   0%|                                                              | 0/5 [00:00<?, ?it/s]





























































































































































Training...:   1%|▎                                                | 500/92767 [07:14<17:27:00,  1.47it/s]






























































































































































Training...:   1%|▌                                              | 1001/92767 [13:19<188:17:28,  7.39s/it]




























































































































































Training...:   2%|▊                                               | 1500/92767 [18:38<17:17:45,  1.47it/s]





























































































































































Training...:   2%|█                                               | 2000/92767 [24:20<17:10:47,  1.47it/s]





























































































































































Training...:   3%|█▎                                              | 2500/92767 [30:02<17:06:01,  1.47it/s]