Error while training object detection model in TF2

While training using TensorFlow Object Detection API from Google Colab I got the following error

tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  assertion failed: [[0.397090524]] [[0.199330032]]
     [[{{node Assert/AssertGuard/else/_25/Assert/AssertGuard/Assert}}]]
     [[MultiDeviceIteratorGetNextFromShard]]
     [[RemoteCall]]
     [[IteratorGetNext]]
  (1) Invalid argument:  assertion failed: [[0.397090524]] [[0.199330032]]
     [[{{node Assert/AssertGuard/else/_25/Assert/AssertGuard/Assert}}]]
     [[MultiDeviceIteratorGetNextFromShard]]
     [[RemoteCall]]
     [[IteratorGetNext]]
     [[Loss/classification_loss_1/write_summary/ReadVariableOp/_48]]
0 successful operations.
0 derived errors ignored. [Op:__inference__dist_train_step_62874]

Function call stack:
_dist_train_step -> _dist_train_step

Can someone please help me'