Make a remark about a serious bug I made previously

vit_gpt2/modeling_flax_gpt2.py

@@ -309,6 +309,7 @@ class FlaxGPT2Block(nn.Module):
 
         if not self.only_self_attn:
             self.encoder_ln = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
+            # [IMPORTANT] Cross attention requires ``causal=False``! This is a bug I made previously.
             self.encoder_attn = FlaxGPT2Attention(config=self.config, dtype=self.dtype, causal=False, self_attn=False)
 
         self.ln_2 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)