cuda crash debugging

app.py

@@ -212,7 +212,7 @@ with gr.Blocks(theme="gradio/soft") as demo:
     with gr.Tab("Journal"):
         gr.Markdown(
             """
-            ## On How Four Crazy Fellows Embarked on Training a U-Net from Scratch in Five Days and Almost Died in the End
+            ## On How Four Crazy Fellows Embarked on Training a U-Net from Scratch in Five Days with JAX and Almost Died in the End
 
             Lorem ipsum dolor sit amet, consectetur adipiscing elit. Mauris vitae varius libero. Nullam laoreet eget sapien quis tristique. Cras odio odio, consequat sed cursus quis, dignissim hendrerit ligula. Curabitur non lorem tellus. Nam bibendum malesuada mi sed faucibus. Sed euismod enim metus, sit amet venenatis elit elementum vel. Duis nec rhoncus tellus, rhoncus auctor justo. Proin id gravida dolor. Sed nulla lectus, finibus non fringilla ac, fermentum in sapien. Cras lobortis est augue, vel posuere justo pretium vitae. Aliquam lorem dolor, condimentum et finibus rutrum, rhoncus eget nunc.
 

requirements.txt

@@ -1,5 +1,5 @@
 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-jax[cuda11_local]
+jax[cuda12_pip]
 flax
 optax
 diffusers