Mixed precision with TPU on google colab

Hello,

On Colab, I don’t have the GPU with minimum required compute capability for mixed precision. So I decided to use TPU and set the policy as ‘mixed_bfloat16’. However, while performing transfer learning using EfficientNetB0, I get the following error. I am not able to find solutions online as well. I am also not sure if this is because I am using TPU or if there is something I am missing. Any help would be great. This is the error:

/usr/local/lib/python3.7/dist-packages/keras/utils/generic_utils.py:497: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.

category=CustomMaskWarning)
Epoch 1/2

InvalidArgumentError Traceback (most recent call last)
in ()
----> 1 hist1 = model.fit(train_data,steps_per_epoch=len(train_data),epochs=2, validation_data=test_data,validation_steps=int(0.15*len(test_data)),callbacks = [callback1,callback2])

6 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
58 ctx.ensure_initialized()
59 tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
—> 60 inputs, attrs, num_outputs)
61 except core._NotOkStatusException as e:
62 if name is not None:

InvalidArgumentError: No OpKernel was registered to support Op ‘DepthwiseConv2dNativeBackpropInput’ used by {{node gradient_tape/model_7/efficientnetb0/block7a_dwconv/depthwise/DepthwiseConv2dNativeBackpropInput}} with these attrs: [dilations=[1, 1, 1, 1], T=DT_BFLOAT16, data_format=“NHWC”, strides=[1, 1, 1, 1], explicit_paddings=[], padding=“SAME”]
Registered devices: [CPU]
Registered kernels:
device=‘XLA_CPU_JIT’; T in [DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_HALF]
device=‘XLA_GPU_JIT’; T in [DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_HALF]
device=‘GPU’; T in [DT_DOUBLE]
device=‘GPU’; T in [DT_FLOAT]
device=‘GPU’; T in [DT_HALF]
device=‘CPU’; T in [DT_DOUBLE]
device=‘CPU’; T in [DT_FLOAT]
device=‘CPU’; T in [DT_HALF]
device=‘GPU’; label=‘cudnn_grouped_convolution’; T in [DT_DOUBLE]
device=‘GPU’; label=‘cudnn_grouped_convolution’; T in [DT_FLOAT]
device=‘GPU’; label=‘cudnn_grouped_convolution’; T in [DT_HALF]