ColabのTPUを使用してTensorflowモデルを学習

Question

前提・実現したいこと

Google Colabに搭載されているTPUを使用してEfficientNetモデルを学習したいです。
しかしながら、下記エラーが発生することにより上手くいきません。
どなたか解決方法をご教授頂ければ幸いです。

発生している問題・エラーメッセージ

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-4-a67e6b797143> in <module>()
     20 print('### START TRAINING ###')
     21 
---> 22 history = model.fit(data_train, steps_per_epoch=TRAIN_STEPS_PER_EPOCH, epochs=EPOCS, validation_data=data_val, validation_steps=VAL_STEPS_PER_EPOCH)
     23 
     24 print('### FINISH TRAINING ###')

5 frames
/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Unable to parse tensor proto

該当のソースコード

import tensorflow as tf
from keras.utils import to_categorical

AUTOTUNE = tf.data.experimental.AUTOTUNE
IMAGE_SIZE = 240
TRAIN_IMAGE_COUNT = len(df_train)
VAL_IMAGE_COUNT = len(df_val)
BATCH_SIZE = 32

def preprocess_image(path, label=None):
  image = tf.io.read_file(path)
  image = tf.image.decode_jpeg(image, channels=3)
  image = tf.image.resize(image, [IMAGE_SIZE, IMAGE_SIZE])
  image = tf.cast(image, tf.float32)/255.0

  return image, label

data_train = tf.data.Dataset.from_tensor_slices((df_train['image_path'], df_train['label']))
data_train = data_train.map(preprocess_image)
data_train = data_train.shuffle(buffer_size=TRAIN_IMAGE_COUNT)
data_train = data_train.repeat()
data_train = data_train.batch(BATCH_SIZE)
data_train = data_train.prefetch(buffer_size=AUTOTUNE)
data_train = data_train.cache()

data_val = tf.data.Dataset.from_tensor_slices((df_val['image_path'], df_val['label']))
data_val = data_val.map(preprocess_image)
data_val = data_val.shuffle(buffer_size=VAL_IMAGE_COUNT)
data_val = data_val.repeat()
data_val = data_val.batch(BATCH_SIZE)
data_val = data_val.prefetch(buffer_size=AUTOTUNE)
data_val = data_val.cache()

%tensorflow_version 2.x
import tensorflow as tf
print("Tensorflow version " + tf.__version__)

try:
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
except ValueError:
  raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.TPUStrategy(tpu)

from tensorflow.keras.applications import EfficientNetB1
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout
from tensorflow.keras import optimizers

TRAIN_STEPS_PER_EPOCH = tf.math.ceil(len(df_train)/BATCH_SIZE).numpy()
VAL_STEPS_PER_EPOCH = tf.math.ceil(len(df_val)/BATCH_SIZE).numpy()
EPOCS = 1

with strategy.scope():
  model = tf.keras.Sequential()

  conv_base = EfficientNetB1(weights='imagenet', input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3), include_top=False)

  model.add(conv_base)
  model.add(GlobalAveragePooling2D())
  model.add(Dense(5, activation='softmax'))

  model.compile(loss='sparse_categorical_crossentropy', optimizer=optimizers.Adagrad(), metrics=['accuracy'])

print('### START TRAINING ###')

history = model.fit(data_train, steps_per_epoch=TRAIN_STEPS_PER_EPOCH, epochs=EPOCS, validation_data=data_val, validation_steps=VAL_STEPS_PER_EPOCH)

print('### FINISH TRAINING ###')

試したこと

正直なところ、初めのdatasetを構築するところの意味があまり理解出来ていなが、下記を試した。
効果は得ることが出来なかった。

image = tf.cast(image, tf.float16)/255.0

補足情報（FW/ツールのバージョンなど）

Google Colab
Tensorflow version 2.4.0