Function create_optimizer

utils/optimizer_utils.py:7–16

(model, args)

def create_optimizer(model, args):
    from torch.nn.parallel import DistributedDataParallel as DDP
    raw_model = model.module if isinstance(model, DDP) else model

    optimizer_name = getattr(args, "optimizer", "adamw")

    if optimizer_name == "muon":
        return _create_muon_split(raw_model, args)
    else:
        return _create_adamw(raw_model, args)


def _create_adamw(model, args):