E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2732] Execution of replica 0 failed

[VachanVY]

I get the below error when training

Training about to start...
Saved checkpoint of step 0.
| Step: 0 || Loss: 1.1787 || LR: 0.000000e+00 || dt: 6917.22ms |
| Step: 1 || Loss: 1.1401 || LR: 1.000000e-05 || dt: 166.68ms |
2024-07-12 21:49:19.329449: E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2732] Execution of replica 0 failed: INVALID_ARGUMENT: Executable expected 69 arguments but got 70

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[36], line 27
     25 grads = jax.tree_util.tree_map(jnp.zeros_like, trainable_vars)
     26 for _ in range(config.num_grad_accumalation_steps):
---> 27     grad, loss, (trainable_vars, non_trainable_vars) = mini_step(
     28         (trainable_vars, non_trainable_vars),
     29         noisy_image_timestep["noisy_images"], noisy_image_timestep["timesteps"],
     30         labels, noise_true, key4,
     31         config.num_grad_accumalation_steps
     32     )
     33     grads = jax.tree_util.tree_map(
     34         lambda g1, g2: jnp.add(g1, g2), grads, grad
     35     ) # sum grads for grad accumation
     36     X_batch, labels = debug_input

ValueError: INVALID_ARGUMENT: Executable expected 69 arguments but got 70

but if restart my jupyter notebook the training starts with no errors and sometimes I get the same error back.
Could you let me know why this happens? The code runs for 2 steps though... so no problem in my code.