How to run tensorboard in colab

[sokol11]

Hi. How do I profile my flax network in Colab (TPU)? I tried running

with jax.profiler.trace(log_dir):
    pi, value = apply_network(train_state.params, obsv, rng)
    value.block_until_ready()

But it doesn't look like it's writing the trace to the log folder. I get the following when I launch tensorboard:

"""
No dashboards are active for the current data set.
Probable causes:

You haven’t written any data to your event files.
TensorBoard can’t find your event files.
If you’re new to using TensorBoard, and want to find out how to add data and set up your event files, check out the README and perhaps the TensorBoard tutorial.
If you think TensorBoard is configured properly, please see the section of the README devoted to missing data problems and consider filing an issue on GitHub.

Last reload: Sep 23, 2024, 10:56:23 PM
Log directory: tmp/log/"
"""

Any advice would be greatly appreciated. Thank you!