Converting a tensor into integer value in a tf.function decorated function?

I'm trying to loop over the tfdataset and run a custom train step on a number of replicas using distributed training. My function is as follows:

@tf.function
def distributedTrainEpoch(self, trainData):

   totalLoss = 0.0
   totalDice = 0.0
   i = 0
   prog = Progbar(self.trainSteps-1)

   for batch in trainData:

      i+=1

      replicaLoss, replicaDice = self.strategy.run(self.trainStep, args=(batch,))
      totalLoss += self.strategy.reduce(tf.distribute.ReduceOp.SUM, replicaLoss, axis=None)
      totalDice += self.strategy.reduce(tf.distribute.ReduceOp.SUM, replicaDice, axis=None)
      prog.update(i)

   return totalLoss, totalDice

The problem here is I'm trying to use the Progress bar from tf.keras.utils.Progbar and when I update it with the integer i, I get a type error problem because i is a tensor since we use @tf.function to execute this code in graph mode rather than eagerly.

I know in tensorflow 1, I could start a session and call i.eval() but I can't do that here. Does anyone know how I can convert a tensor to it's value under a tf.function?

I could obviously remove the loop from this function but I don't want to do that.

Thanks,