iree: Error when compiling the Bert large model

Describe the bug Hello, I have an error trying to compile the Bert large model. The error is as below

/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/distribute/merge_call_interim.py:51:0: error: 'util.global.store' op global type mismatch; global __iree_flow___sm_node20__optimizer.iter is 'tensor<i64>' but store is 'tensor<i32>'
/home/vivian/.local/lib/python3.9/site-packages/keras/optimizer_v2/optimizer_v2.py:671:0: note: called from
/home/vivian/iree-samples/ModelCompiler/nlp_models/bert_small_gen.py:60:0: note: called from
/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/framework/func_graph.py:1136:0: note: called from
/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/eager/function.py:3831:0: note: called from
/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/eager/def_function.py:677:0: note: called from
/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/framework/func_graph.py:1161:0: note: called from
/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/eager/function.py:3130:0: note: called from
/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/eager/function.py:3292:0: note: called from
/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/eager/function.py:3036:0: note: called from
/home/vivian/.local/lib/python3.9/site-packages/tensorflow/python/distribute/merge_call_interim.py:51:0: note: see current operation: "util.global.store"(%2036) {global = @__iree_flow___sm_node20__optimizer.iter} : (tensor<i32>) -> ()
compilation from source to vm failed

To Reproduce

1. Generate the Bert large model using the script from here https://github.com/google/iree-samples/blob/main/ModelCompiler/nlp_models/bert_large_gen.py
2. Run command with CUDA NVIDIA V100 ./iree-compile -iree-input-type=mhlo -iree-mlir-to-vm-bytecode-module -iree-hal-target-backends=cuda -iree-cuda-llvm-target-arch=sm_70 -iree-hal-cuda-disable-loop-nounroll-wa /home/vivian/iree-samples/ModelCompiler/nlp_models/model.mlir -o model.vmfb

Additional context

1. IREE was built with this commit cbea6b2670f56ba20134408b222d41c911f36cf8
2. IREE pip packages 20220325.87