Unsupported data type for TPU String

Getting this error when running baselines/cifar10/deterministic.py

tensorflow/compiler/jit/xla_compilation_cache.cc:334] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process.I1017 19:47:33.579260 140642886925376 deterministic.py:260] [deterministic.py:260] Model input shape: (None, 32, 32, 3)
I1017 19:47:33.579456 140642886925376 deterministic.py:261] [deterministic.py:261] Model output shape: (None, 10)
I1017 19:47:33.581461 140642886925376 deterministic.py:262] [deterministic.py:262] Model number of weights: 36497146
I1017 19:47:34.059113 140642886925376 deterministic.py:457] [deterministic.py:457] Starting to run epoch: 0
I1017 19:47:34.920352 140642886925376 tpu.py:1376] [tpu.py:1376] TPU has inputs with dynamic shapes: [, , , , , ]
I1017 19:47:40.580166 140642886925376 tpu.py:1376] [tpu.py:1376] TPU has inputs with dynamic shapes: [, <
tf.Tensor 'while/cond_8/Identity:0' shape=(None,) dtype=int64>, , , , ]
2021-10-17 19:47:44.299236: I tensorflow/core/tpu/graph_rewrite/encapsulate_tpu_computations_pass.cc:263] Subgraph fingerprint:1696516458553054063
Traceback (most recent call last):
File "baselines/cifar/deterministic.py", line 557, in
app.run(main)
File "/usr/local/lib/python3.8/dist-packages/absl/app.py", line 303, in run
_run_main(main, args)
File "/usr/local/lib/python3.8/dist-packages/absl/app.py", line 251, in _run_main
sys.exit(main(argv))
File "baselines/cifar/deterministic.py", line 462, in main
train_step(train_iterator)
File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/def_function.py", line 885, in __call__
result = self._call(*args, **kwds)
File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/def_function.py", line 946, in _call
return self._stateless_fn(*args, **kwds)
File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/function.py", line 3032, in __call__
return graph_function._call_flat(
File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/function.py", line 1956, in _call_flat
return self._build_call_outputs(self._inference_function.call(
File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/function.py", line 591, in call
outputs = execute.execute(
File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/execute.py", line 59, in quick_execute
tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: Unsupported data type for TPU: string, caused by output while/body/_1/while/Pad_3:0 [Op:__i
nference_train_step_26492]`