Tensorboard profile failed to capture when xla enabled

TensorFlow General Discussion
, , ,
1

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)

profile
profile1465×585 63.8 KB

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