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.
@tf.function(jit_compile=True) 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.