def _log_lamb_rs(optimizer: Optimizer, event_writer, token_count: int):
    """Log a histogram of trust ratio scalars in across layers."""
    results = collections.defaultdict(list)
    for group in optimizer.param_groups:
        for p in group["params"]:
            state = optimizer.state[p]
            for i in ("weight_norm", "adam_norm", "trust_ratio"):
                if i in state:
                    results[i].append(state[i])

    for k, v in results.items():
        event_writer.add_histogram(f"lamb/{k}", torch.tensor(v), token_count)
Log a histogram of trust ratio scalars in across layers.
catalyst/contrib/optimizers/lamb.py