Fix run_tpu_test (#14641)

Commit

1 year ago

Fix run_tpu_test (#14641) Signed-off-by: <ricliu@google.com> Signed-off-by: Richard Liu <ricliu@google.com>