[mlgo][nfc] regalloc test model generator: prep for TFLite

Casting operator to make TFLite happy.

Reviewed By: yundiqian

Differential Revision: https://reviews.llvm.org/D131584

Author: mtrofin

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

@@ -46,7 +46,8 @@ def build_mock_model(path):
   module.var = tf.Variable(0, dtype=tf.int64)
 
   def action(*inputs):
-    result = tf.math.argmax(tf.cast(inputs[0]['mask'], tf.int32), 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 = {