Update requirements.txt

requirements.txt

@@ -3,6 +3,7 @@ absl-py>=0.12.0
 chex>=0.0.7
 clu>=0.0.3
 einops>=0.3.0
+# flax==0.6.9
 git+https://github.com/google/flaxformer
 # jax[cuda11_cudnn86]>=0.4.2
 #--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html