fix mlgo regalloc test model generation for tflite

[yundiqian]

Summary

To move from TF C API to TFLite, we found that the argmax op in TFLite does not work for int64 inputs, so cast the int64 inputs to int32 inputs to make TFLite argmax op work

llvm/lib/Analysis/models/gen-regalloc-eviction-test-model.py

 def build_mock_model(path):
   """Build and save the mock model with the given signature."""
   module = tf.Module()
   # We have to set this useless variable in order for the TF C API to correctly
   # intake it
   module.var = tf.Variable(0, dtype=tf.int64)
 
   def action(*inputs):
-    result = tf.math.argmax(inputs[0]['mask'], axis=-1) + module.var
+    result = tf.math.argmax(tf.cast(inputs[0]['mask'], tf.int32), axis=-1) + module.var
     return {POLICY_DECISION_LABEL: result}
   module.action = tf.function()(action)
   action = {
       'action': module.action.get_concrete_function(get_input_signature())
   }
   tf.saved_model.save(module, path, signatures=action)
   output_spec_path = get_output_spec_path(path)
   with open(output_spec_path, 'w') as f: