#用8bit adam优化器节省显存

# Initialize the optimizer
if cfg.solver.use_8bit_adam:
    try:
        import bitsandbytes as bnb
    except ImportError:
        raise ImportError(
            "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
        )

    optimizer_cls = bnb.optim.AdamW8bit
else:
    optimizer_cls = torch.optim.AdamW

optimizer = optimizer_cls(
    trainable_params,
    lr=learning_rate,
    betas=(cfg.solver.adam_beta1, cfg.solver.adam_beta2),
    weight_decay=cfg.solver.adam_weight_decay,
    eps=cfg.solver.adam_epsilon,
)