fix gelu block.

Author: jfacevedo-google

src/maxdiffusion/configs/base_wan_14b.yml

@@ -102,6 +102,7 @@ base_output_directory: ""
 
 # Hardware
 hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
+skip_jax_distributed_system: False
 
 # Parallelism
 mesh_axes: ['data', 'fsdp', 'tensor']