JAX / FLAX Training speed suddenly slows down

[jgnooo]

Hi,

I'm customizing tensorf-jax.
But, when training, the training step suddenly slow

Iteration 001/100 --> elapsed: 23.5022 seconds.
Iteration 002/100 --> elapsed: 0.0013 seconds.
Iteration 003/100 --> elapsed: 0.0018 seconds.
Iteration 004/100 --> elapsed: 0.0011 seconds.
Iteration 005/100 --> elapsed: 0.0010 seconds.
Iteration 006/100 --> elapsed: 0.0010 seconds.
Iteration 007/100 --> elapsed: 0.0010 seconds.
Iteration 008/100 --> elapsed: 0.0010 seconds.
Iteration 009/100 --> elapsed: 0.0010 seconds.
Iteration 010/100 --> elapsed: 0.0011 seconds.
Iteration 011/100 --> elapsed: 0.0010 seconds.
Iteration 012/100 --> elapsed: 0.0010 seconds.
Iteration 013/100 --> elapsed: 0.0010 seconds.
Iteration 014/100 --> elapsed: 0.0010 seconds.
Iteration 015/100 --> elapsed: 0.0010 seconds.
Iteration 016/100 --> elapsed: 0.0010 seconds.
Iteration 017/100 --> elapsed: 0.0010 seconds.
Iteration 018/100 --> elapsed: 0.0010 seconds.
Iteration 019/100 --> elapsed: 0.0010 seconds.
Iteration 020/100 --> elapsed: 0.0010 seconds.
Iteration 021/100 --> elapsed: 0.0215 seconds.
Iteration 022/100 --> elapsed: 0.1074 seconds.
Iteration 023/100 --> elapsed: 0.1073 seconds.
Iteration 024/100 --> elapsed: 0.1074 seconds.
Iteration 025/100 --> elapsed: 0.1074 seconds.
Iteration 026/100 --> elapsed: 0.1076 seconds.
Iteration 027/100 --> elapsed: 0.1076 seconds.
Iteration 028/100 --> elapsed: 0.1074 seconds.
Iteration 029/100 --> elapsed: 0.1085 seconds.
Iteration 030/100 --> elapsed: 0.1058 seconds.
Iteration 031/100 --> elapsed: 0.1071 seconds.
Iteration 032/100 --> elapsed: 0.1076 seconds.
Iteration 033/100 --> elapsed: 0.1073 seconds.
Iteration 034/100 --> elapsed: 0.1069 seconds.
Iteration 035/100 --> elapsed: 0.1073 seconds.

Do you know why?, and What do I have to do?