iree: SGD Not working with BERT

Hey guys, I am trying to compile (w/ iree.compiler.tf.compile_module) BERT. Inference/predict seems to work fine. But when trying to compile learn/training, I got this following error:

/path/to/python3.8/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:714:0: error: could not lower resource op to flow: tf.ResourceScatterAdd

To reproduce:

pip install gin-config
git clone https://github.com/raikonenfnu/ModelCompiler 
export PYTHONPATH=${PYTHONPATH}:$PWD/ModelCompiler/third_party/tf_models
cd ModelCompiler/nlp_gen
python bert_gen.py

IREE SNAPSHOT VERSION - releases/tag/snapshot-20210505.259

Thanks in advance guys 😃