fix minor typo in shape

Author: kmario23

pdebench/data_gen/data_gen_NLE/utils.py

@@ -1140,7 +1140,7 @@ def __create_3DRand_init(u, d0, T0, delD, delP, keys):
                         continue
                     # random phase
                     key, subkey = random.split(key)
-                    phs = 2. * jnp.pi * random.uniform(key, shape=[3])  # (vi, k)
+                    phs = 2. * jnp.pi * random.uniform(key, shape=[5])  # (vi, k)
 
                     uk = 1./jnp.sqrt(kx**2 + ky**2 + kz**2)
                     kdx = kx * xc[:,None,None] + ky * yc[None,:,None] + kz * zc[None,None,:]