import torch


def setup_optimizer_and_scheduler(param_lr_maps, base_lr, epochs, steps_per_epoch):
    """Create a PyTorch AdamW optimizer and OneCycleLR scheduler with `param_lr_maps` parameter mapping,
       with base LR `base_lr`, for training for `epochs` epochs, with `steps_per_epoch` iterations
       per epoch."""
    optimizer = torch.optim.AdamW(param_lr_maps, lr=base_lr)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, base_lr, epochs=epochs, steps_per_epoch=steps_per_epoch
    )
    return optimizer, scheduler