Program suddenly hangs

[fmscole]

Has anyone ever encountered a situation where a program that was originally running well suddenly hangs without any error messages and does not end or run. And when to hang is random.

What is the most likely problem in this situation? Memory is absolutely sufficient.

A friend on the Internet said that GPU was waiting for data, but the data did not enter GPU, so how to solve it?

https://github.com/fmscole/jax_example/blob/test/lstm/main.py

@jax.jit
def predict(state, variables,image_i):  
  logits= state.apply_fn({'params': state.params},image_i,is_training=False)
  return logits

System information：

• Win11(64G)+wsl2+ Linux Ubuntu 18.04
• Flax(0.5.2), jax(0.3.14), jaxlib(0.3.14+cuda11.cudnn82)
• Python 3.8.8:
• GPU 3080lap 16G:
• CUDA version: 11.5

[cgarciae]

Hey @fmscole! I don't think this is Flax related so I turned the issue into a discussion.

I've had stuff like this happen either because of the jax version or even the cuda version, its hard to know. Test with different JAX and cuda versions if possible.