Tensorboard profile failed to capture when xla enabled

I am trying to capture gpu utilization during model inference using tensorboard profile. This works fine as long as when xla is not enabled(i.e when jit_compile=False). When xla is enabled tensorboard profile says “No step marker observed and hence the step time is unknown.”.

This is how inference is done with xla enabled.

    def do_inference(self, layer, img):
        return layer(img)

I couldn’t find what is the issue here. Let me know if anyone has any clue.