tb_logger = TensorboardLogger(log_dir="experiments/tb_logs")

tb_logger.attach_output_handler(
    trainer,
    event_name=Events.ITERATION_COMPLETED,
    tag="training",
    state_attrs=["alpha", "beta"]  # <->  trainer.state.alpha and trainer.state.beta
)

tb_logger = TensorboardLogger(log_dir="experiments/tb_logs")

def trainer_state_logger(state_attrs, tag=None):

    def wrapper(engine, logger, event_name):
        assert engine == trainer

        global_step = engine.state.get_event_attrib_value(event_name)
        tag_prefix = f"{tag}/" if tag else ""
        for name in state_attrs:
            value = getattr(engine, name, None)
            if value is None:
                continue
            logger.writer.add_scalar(f"{tag_prefix}{name}", value, global_step)
    
tb_logger.attach(
    trainer,
    event_name=Events.ITERATION_COMPLETED,
    log_handler=trainer_state_logger(["alpha", "beta"], tag="training")
)