renoved extra debug

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -225,7 +225,6 @@ def _tpu_flash_attention(
     attention_mask: jax.Array = None,
 ) -> jax.Array:
   """TPU Flash Attention"""
-  jax.debug.print("USing FLASH ATTENTION")
 
   q_max_block_size = 1024 if dtype == jnp.bfloat16 else 512
   # This is the case for cross-attn.
@@ -445,7 +444,6 @@ def _apply_attention_dot(
     float32_qk_product: bool,
     use_memory_efficient_attention: bool,
 ):
-  jax.debug.print("Using DOT PRODUCT ATTENTION")
   """Apply Attention."""
   if split_head_dim:
     b = key.shape[0]