Code update

PiperOrigin-RevId: 669491071

swirl_dynamics/lib/diffusion/unets.py

@@ -573,7 +573,7 @@ def __call__(self, emb: Array, cond: dict[str, Array], is_training: bool):
         dtype=self.dtype,
         param_dtype=self.param_dtype,
     )(value.reshape(b, -1))
-    value = nn.swish(value)
+    value = nn.swish(nn.LayerNorm()(value))
 
     # Concatenate the noise and conditional embedding.
     emb = jnp.concatenate([emb, value], axis=-1)