Tensorboard SummaryWriter not working

[trechberger]

Hi!

Maybe I'm missing something, but is there a bug in flax.metrics.tensorboard.SummaryWriter?

The following minimal example is producing a TypeError:

from flax.metrics.tensorboard import SummaryWriter
summary_writer = SummaryWriter("./summary-writer-test/")
summary_writer.scalar("test", 1, 1)

->
TypeError: __exit__() missing 3 required positional arguments: 'type', 'value', and 'traceback'

The error suggests that in flax/metrics/tensorboard.py:70, the arguments are missing and should be provided (see here).

Thanks for any feedback!